@triton.jit
def _fwd_kernel_inner(acc, l_i, m_i, q, Q, k_block_col_idx, layout_col_ptr,
    layout_col_stride_h, layout_col_stride_m, k_ptrs, v_ptrs, off_h, offs_m,
    offs_n, offs_d, stride_kt, stride_vt, sm_scale, k_seqlen, past_len,
    LAST_K_BLOCK: tl.constexpr, BLOCK_M_LOADING: tl.constexpr, BLOCK_N: tl.
    constexpr, D_HEAD: tl.constexpr, EVEN_D: tl.constexpr, M_LT_N: tl.constexpr
    ):
    k_block_id = tl.load(layout_col_ptr + off_h * layout_col_stride_h + 
        k_block_col_idx * layout_col_stride_m).to(tl.int32)
    start_n = k_block_id * BLOCK_N
    if LAST_K_BLOCK:
        if EVEN_D:
            k = tl.load(k_ptrs + start_n * stride_kt, mask=offs_n[None, :] +
                start_n < k_seqlen)
        else:
            k = tl.load(k_ptrs + start_n * stride_kt, mask=(offs_n[None, :] +
                start_n < k_seqlen) & (offs_d[:, None] < D_HEAD))
    elif EVEN_D:
        k = tl.load(k_ptrs + start_n * stride_kt)
    else:
        k = tl.load(k_ptrs + start_n * stride_kt, mask=offs_d[:, None] < D_HEAD
            )
    qk = tl.zeros([BLOCK_M_LOADING, BLOCK_N], dtype=tl.float32)
    qk += tl.dot(q, k)
    qk *= sm_scale
    if LAST_K_BLOCK | M_LT_N:
        qk += tl.where(offs_m[:, None] + past_len >= start_n + offs_n[None,
            :], 0, float('-inf'))
    m_ij = tl.maximum(m_i, tl.max(qk, 1))
    p = tl.math.exp2(qk - m_ij[:, None])
    l_ij = tl.sum(p, 1)
    alpha = tl.math.exp2(m_i - m_ij)
    acc = acc * alpha[:, None]
    m_i = m_ij
    l_i = l_i * alpha + l_ij
    p = p.to(Q.dtype.element_ty)
    if LAST_K_BLOCK:
        if EVEN_D:
            v = tl.load(v_ptrs + start_n * stride_vt, mask=offs_n[:, None] +
                start_n < k_seqlen)
        else:
            v = tl.load(v_ptrs + start_n * stride_vt, mask=(offs_n[:, None] +
                start_n < k_seqlen) & (offs_d[None, :] < D_HEAD))
    elif EVEN_D:
        v = tl.load(v_ptrs + start_n * stride_vt)
    else:
        v = tl.load(v_ptrs + start_n * stride_vt, mask=offs_d[None, :] < D_HEAD
            )
    acc += tl.dot(p, v)
    return acc, l_i, m_i
