@triton.jit
def _vq_attn_fwd_inner(acc, l_i, m_i, q, sm_scale, K_block_ptr,
    K_VQC_block_ptr, VVQI_block_ptr, V_VQ, stride_vvq_n, CODEBOOK_SIZE: tl.
    constexpr, BLOCK_HEADDIM: tl.constexpr, BLOCK_N: tl.constexpr):
    for _ in range(0, CODEBOOK_SIZE, BLOCK_N):
        k = tl.load(K_block_ptr)
        qk = tl.dot(q, k) * (sm_scale * RCP_LN2)
        k_vq_cnt = tl.load(K_VQC_block_ptr)
        mask = k_vq_cnt != 0.0
        qk = tl.where(mask[None, :], qk, NEGINF)
        m_i_new = tl.maximum(m_i, tl.max(qk, 1))
        alpha = tl.math.exp2(m_i - m_i_new)
        p = tl.math.exp2(qk - m_i_new[:, None])
        acc *= alpha[:, None]
        v_vq_index = tl.load(VVQI_block_ptr)
        v_ptr = V_VQ + (v_vq_index[:, None] * stride_vvq_n + tl.arange(0,
            BLOCK_HEADDIM)[None, :])
        v = tl.load(v_ptr)
        acc += tl.dot(p.to(v.dtype), v)
        l_i = l_i * alpha + tl.sum(p.to(v.dtype) * k_vq_cnt[None, :], axis=1)
        m_i = m_i_new
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        K_VQC_block_ptr = tl.advance(K_VQC_block_ptr, (BLOCK_N,))
        VVQI_block_ptr = tl.advance(VVQI_block_ptr, (BLOCK_N,))
    return acc, l_i, m_i
