@triton.autotune(configs=[triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64,
    'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1, num_warps=8),
    triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'waves_per_eu': 2,
    'PRE_LOAD_V': False}, num_stages=1, num_warps=4), triton.Config({
    'BLOCK_M': 256, 'BLOCK_N': 128, 'waves_per_eu': 2, 'PRE_LOAD_V': False},
    num_stages=1, num_warps=8), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 
    64, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1, num_warps=4),
    triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 3,
    'PRE_LOAD_V': True}, num_stages=1, num_warps=4), triton.Config({
    'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 3, 'PRE_LOAD_V': False},
    num_stages=1, num_warps=4), triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64,
    'waves_per_eu': 4, 'PRE_LOAD_V': False}, num_stages=1, num_warps=8),
    triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'waves_per_eu': 4,
    'PRE_LOAD_V': False}, num_stages=1, num_warps=8), triton.Config({
    'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 1, 'PRE_LOAD_V': False},
    num_stages=1, num_warps=4)], key=['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL']
    )
@triton.jit
def attn_fwd(Q, K, V, bias, sm_scale, L, Out, stride_qz, stride_qh,
    stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk,
    stride_vz, stride_vh, stride_vk, stride_vn, stride_oz, stride_oh,
    stride_om, stride_on, stride_bz, stride_bh, stride_bm, stride_bn,
    cu_seqlens_q, cu_seqlens_k, dropout_p, philox_seed, philox_offset_base,
    encoded_softmax, HQ: tl.constexpr, HK: tl.constexpr,
    ACTUAL_BLOCK_DMODEL: tl.constexpr, MAX_SEQLENS_Q: tl.constexpr,
    MAX_SEQLENS_K: tl.constexpr, VARLEN: tl.constexpr, IS_CAUSAL: tl.
    constexpr, BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N:
    tl.constexpr, PRE_LOAD_V: tl.constexpr, BIAS_TYPE: tl.constexpr,
    ENABLE_DROPOUT: tl.constexpr, RETURN_ENCODED_SOFTMAX: tl.constexpr):
    start_m = tl.program_id(0)
    off_h_q = tl.program_id(1)
    off_z = tl.program_id(2)
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    if VARLEN:
        cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z)
        cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1)
        seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start
        if start_m * BLOCK_M > seqlen_q:
            return
        cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z)
        cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1)
        seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start
    else:
        cu_seqlens_q_start = 0
        cu_seqlens_k_start = 0
        seqlen_q = MAX_SEQLENS_Q
        seqlen_k = MAX_SEQLENS_K
    n_blocks = cdiv_fn(seqlen_k, BLOCK_N)
    if IS_CAUSAL:
        n_blocks_seqlen = cdiv_fn((start_m + 1) * BLOCK_M + seqlen_k -
            seqlen_q, BLOCK_N)
        n_blocks = min(n_blocks, n_blocks_seqlen)
        if n_blocks <= 0:
            o_offset = (off_z * stride_oz + cu_seqlens_q_start * stride_om +
                off_h_q * stride_oh)
            O_block_ptr = tl.make_block_ptr(base=Out + o_offset, shape=(
                seqlen_q, BLOCK_DMODEL), strides=(stride_om, stride_on),
                offsets=(start_m * BLOCK_M, 0), block_shape=(BLOCK_M,
                BLOCK_DMODEL), order=(1, 0))
            acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=Out.type.element_ty)
            return
    GROUP_SIZE: tl.constexpr = HQ // HK
    off_h_k = off_h_q // GROUP_SIZE if GROUP_SIZE != 1 else off_h_q
    n_extra_tokens = 0
    if seqlen_k < BLOCK_N:
        n_extra_tokens = BLOCK_N - seqlen_k
    elif seqlen_k % BLOCK_N:
        n_extra_tokens = seqlen_k % BLOCK_N
    padded_head = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL
    q_offset = (off_z * stride_qz + off_h_q * stride_qh + 
        cu_seqlens_q_start * stride_qm)
    Q_block_ptr = tl.make_block_ptr(base=Q + q_offset, shape=(seqlen_q,
        ACTUAL_BLOCK_DMODEL), strides=(stride_qm, stride_qk), offsets=(
        start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(
        1, 0))
    k_offset = (off_z * stride_kz + off_h_k * stride_kh + 
        cu_seqlens_k_start * stride_kn)
    K_block_ptr = tl.make_block_ptr(base=K + k_offset, shape=(
        ACTUAL_BLOCK_DMODEL, seqlen_k), strides=(stride_kk, stride_kn),
        offsets=(0, 0), block_shape=(BLOCK_DMODEL, BLOCK_N), order=(0, 1))
    v_offset = (off_z * stride_vz + off_h_k * stride_vh + 
        cu_seqlens_k_start * stride_vk)
    V_block_ptr = tl.make_block_ptr(base=V + v_offset, shape=(seqlen_k,
        ACTUAL_BLOCK_DMODEL), strides=(stride_vk, stride_vn), offsets=(0, 0
        ), block_shape=(BLOCK_N, BLOCK_DMODEL), order=(1, 0))
    if BIAS_TYPE != 0:
        bias_ptr = tl.make_block_ptr(base=bias + off_h_q * stride_bh, shape
            =(seqlen_q, seqlen_k), strides=(stride_bm, stride_bn), offsets=
            (start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_N), order=(
            1, 0))
    else:
        bias_ptr = None
    if ENABLE_DROPOUT:
        batch_philox_offset = philox_offset_base + (off_z * HQ + off_h_q
            ) * seqlen_q * seqlen_k
    else:
        batch_philox_offset = 0
    if RETURN_ENCODED_SOFTMAX:
        encoded_softmax_block_ptr = tl.make_block_ptr(base=encoded_softmax +
            off_h_q * seqlen_q * seqlen_k, shape=(seqlen_q, seqlen_k),
            strides=(seqlen_k, 1), offsets=(start_m * BLOCK_M, 0),
            block_shape=(BLOCK_M, BLOCK_N), order=(1, 0))
    else:
        encoded_softmax_block_ptr = 0
    m_i = tl.full([BLOCK_M], float('-inf'), dtype=tl.float32)
    l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
    qk_scale = sm_scale * 1.44269504089
    q = load_fn(Q_block_ptr, True, padded_head, 'zero')
    q = (q * qk_scale).to(Q_block_ptr.type.element_ty)
    padded_block_k = n_extra_tokens != 0
    is_modulo_mn = not padded_block_k and seqlen_q % BLOCK_M == 0
    if IS_CAUSAL:
        masked_blocks = BLOCK_M // BLOCK_N + (not is_modulo_mn)
    else:
        masked_blocks = padded_block_k
    masked_blocks = min(masked_blocks, n_blocks)
    n_full_blocks = n_blocks - masked_blocks
    block_min = 0
    block_max = n_blocks * BLOCK_N
    if n_full_blocks > 0:
        block_max = (n_blocks - masked_blocks) * BLOCK_N
        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr,
            V_block_ptr, start_m, seqlen_k, dropout_p, philox_seed,
            batch_philox_offset, encoded_softmax_block_ptr, block_min,
            block_max, 0, 0, 0, bias_ptr, False, BLOCK_M, BLOCK_DMODEL,
            BLOCK_N, offs_m, offs_n, PRE_LOAD_V, False, ENABLE_DROPOUT,
            RETURN_ENCODED_SOFTMAX, padded_head)
        block_min = block_max
        block_max = n_blocks * BLOCK_N
    tl.debug_barrier()
    if masked_blocks > 0:
        offs_n_causal = offs_n + (seqlen_q - seqlen_k) if IS_CAUSAL else 0
        K_block_ptr = tl.advance(K_block_ptr, (0, n_full_blocks * BLOCK_N))
        V_block_ptr = tl.advance(V_block_ptr, (n_full_blocks * BLOCK_N, 0))
        if bias_ptr is not None:
            bias_ptr = tl.advance(bias_ptr, (0, n_full_blocks * BLOCK_N))
        if RETURN_ENCODED_SOFTMAX:
            encoded_softmax_block_ptr = tl.advance(encoded_softmax_block_ptr,
                (0, n_full_blocks))
        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr,
            V_block_ptr, start_m, seqlen_k, dropout_p, philox_seed,
            batch_philox_offset, encoded_softmax_block_ptr, block_min,
            block_max, offs_n_causal, masked_blocks, n_extra_tokens,
            bias_ptr, IS_CAUSAL, BLOCK_M, BLOCK_DMODEL, BLOCK_N, offs_m,
            offs_n, PRE_LOAD_V, True, ENABLE_DROPOUT,
            RETURN_ENCODED_SOFTMAX, padded_head)
    acc = acc / l_i[:, None]
    if ENABLE_DROPOUT:
        acc = acc / (1 - dropout_p)
    end_m_idx = (start_m + 1) * BLOCK_M
    start_m_idx = start_m * BLOCK_M
    causal_start_idx = seqlen_q - seqlen_k
    acc = acc.to(Out.type.element_ty)
    if IS_CAUSAL:
        if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx:
            out_mask_boundary = tl.full((BLOCK_DMODEL,), causal_start_idx,
                dtype=tl.int32)
            mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M)
            out_ptrs_mask = mask_m_offsets[:, None] >= out_mask_boundary[
                None, :]
            z = 0.0
            acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty))
    o_offset = (off_z * stride_oz + cu_seqlens_q_start * stride_om + 
        off_h_q * stride_oh)
    O_block_ptr = tl.make_block_ptr(base=Out + o_offset, shape=(seqlen_q,
        ACTUAL_BLOCK_DMODEL), strides=(stride_om, stride_on), offsets=(
        start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(
        1, 0))
    tl.store(O_block_ptr, acc, boundary_check=(0, 1))
