@triton.jit
def attn_fwd(Q, K, V, sm_scale, M, 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, seqlen_q, seqlen_k, dropout_p, philox_seed,
    philox_offset_base, encoded_softmax, STAGE: tl.constexpr, BLOCK_M: tl.
    constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
    pre_load_v: tl.constexpr, ENABLE_DROPOUT: tl.constexpr,
    RETURN_ENCODED_SOFTMAX: tl.constexpr):
    start_m = tl.program_id(0)
    off_h = tl.program_id(1)
    off_z = tl.program_id(2)
    num_h = tl.num_programs(1)
    num_z = tl.num_programs(2)
    q_offset = off_h * stride_qh + off_z * stride_qz
    Q_block_ptr = tl.make_block_ptr(base=Q + q_offset, shape=(seqlen_q,
        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_h * stride_kh + off_z * stride_kz
    K_block_ptr = tl.make_block_ptr(base=K + k_offset, shape=(BLOCK_DMODEL,
        seqlen_k), strides=(stride_kk, stride_kn), offsets=(0, 0),
        block_shape=(BLOCK_DMODEL, BLOCK_N), order=(0, 1))
    v_offset = off_h * stride_vh + off_z * stride_vz
    V_block_ptr = tl.make_block_ptr(base=V + v_offset, shape=(seqlen_k,
        BLOCK_DMODEL), strides=(stride_vk, stride_vn), offsets=(0, 0),
        block_shape=(BLOCK_N, BLOCK_DMODEL), order=(1, 0))
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
    qk_scale = sm_scale * 1.44269504
    q = tl.load(Q_block_ptr)
    q = (q * qk_scale).to(Q_block_ptr.type.element_ty)
    off_zh = off_z * num_h + off_h * 1
    if ENABLE_DROPOUT:
        batch_philox_offset = philox_offset_base + off_zh * 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_zh * 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
    if STAGE & 1:
        acc, l_i, m_i = attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr,
            V_block_ptr, start_m, seqlen_q, seqlen_k, dropout_p,
            philox_seed, batch_philox_offset, encoded_softmax_block_ptr,
            BLOCK_M, BLOCK_DMODEL, BLOCK_N, 4 - STAGE, offs_m, offs_n,
            pre_load_v, ENABLE_DROPOUT, RETURN_ENCODED_SOFTMAX)
    if STAGE & 2:
        tl.debug_barrier()
        acc, l_i, m_i = attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr,
            V_block_ptr, start_m, seqlen_q, seqlen_k, dropout_p,
            philox_seed, batch_philox_offset, encoded_softmax_block_ptr,
            BLOCK_M, BLOCK_DMODEL, BLOCK_N, 2, offs_m, offs_n, pre_load_v,
            ENABLE_DROPOUT, RETURN_ENCODED_SOFTMAX)
    acc = acc / l_i[:, None]
    if ENABLE_DROPOUT:
        acc = acc / (1 - dropout_p)
    m_ptrs = M + off_zh * seqlen_q + offs_m
    tl.store(m_ptrs, m_i + tl.math.log2(l_i))
    o_offset = off_h * stride_oh + off_z * stride_oz
    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))
    tl.store(O_block_ptr, acc.to(Out.type.element_ty))
