@triton.jit
def 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: tl.constexpr, BLOCK_DMODEL: tl.
    constexpr, BLOCK_N: tl.constexpr, STAGE: tl.constexpr, offs_m: tl.
    constexpr, offs_n: tl.constexpr, pre_load_v: tl.constexpr,
    ENABLE_DROPOUT: tl.constexpr, RETURN_ENCODED_SOFTMAX: tl.constexpr):
    if STAGE == 1:
        lo, hi = 0, min(seqlen_k, start_m * BLOCK_M)
    elif STAGE == 2:
        lo, hi = start_m * BLOCK_M, min(seqlen_k, start_m * BLOCK_M + BLOCK_M)
        lo = tl.multiple_of(lo, BLOCK_M)
        K_block_ptr = tl.advance(K_block_ptr, (0, lo))
        V_block_ptr = tl.advance(V_block_ptr, (lo, 0))
        if RETURN_ENCODED_SOFTMAX:
            encoded_softmax_block_ptr = tl.advance(encoded_softmax_block_ptr,
                (0, lo))
    else:
        lo, hi = 0, seqlen_k
    for start_n in range(lo, hi, BLOCK_N):
        if STAGE == 1 or STAGE == 3:
            start_n = tl.multiple_of(start_n, BLOCK_N)
        k = tl.load(K_block_ptr)
        if pre_load_v:
            v = tl.load(V_block_ptr)
        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        if STAGE == 2:
            mask = offs_m[:, None] >= start_n + offs_n[None, :]
            qk = tl.where(mask, qk, float('-inf'))
        if BLOCK_M == 1:
            qk += tl.sum(tl.view(q, [BLOCK_DMODEL]) * tl.view(k, [
                BLOCK_DMODEL]))
        else:
            qk += tl.dot(q, k)
        m_ij = tl.maximum(m_i, tl.max(qk, 1))
        qk = qk - m_ij[:, None]
        p = tl.math.exp2(qk)
        l_ij = tl.sum(p, 1)
        if ENABLE_DROPOUT:
            philox_offset = (batch_philox_offset + start_m * BLOCK_M *
                seqlen_k + start_n)
            keep = dropout_mask(philox_seed, philox_offset, dropout_p,
                BLOCK_M, BLOCK_N, seqlen_k)
            if RETURN_ENCODED_SOFTMAX:
                tl.store(encoded_softmax_block_ptr, tl.where(keep, p, -p).
                    to(encoded_softmax_block_ptr.type.element_ty))
            p = tl.where(keep, p, 0.0)
        elif RETURN_ENCODED_SOFTMAX:
            tl.store(encoded_softmax_block_ptr, p.to(
                encoded_softmax_block_ptr.type.element_ty))
        alpha = tl.math.exp2(m_i - m_ij)
        acc = acc * alpha[:, None]
        if not pre_load_v:
            v = tl.load(V_block_ptr)
        """
        if ENABLE_DROPOUT:
            v = (v / (1.0 - dropout_p)).to(V_block_ptr.type.element_ty)
        """
        l_i = l_i * alpha + l_ij
        m_i = m_ij
        if BLOCK_M == 1:
            acc += tl.view(p.to(V_block_ptr.type.element_ty), [1]) * tl.view(v,
                [BLOCK_DMODEL])
        else:
            acc += tl.dot(p.to(V_block_ptr.type.element_ty), v)
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        if RETURN_ENCODED_SOFTMAX:
            encoded_softmax_block_ptr = tl.advance(encoded_softmax_block_ptr,
                (0, BLOCK_N))
    return acc, l_i, m_i
