@triton.heuristics({'EVEN_M': lambda args: args['seqlen_q'] % args[
    'BLOCK_M'] == 0})
@triton.jit
def _vq_fwd_kernel(Q, K_VQ, K_VQ_CNT, V_VQ, V_VQ_INDEX, Out, L,
    softmax_scale, stride_q_b, stride_q_h, stride_q_m, stride_kvq_h,
    stride_kvq_c, stride_kvqc_b, stride_kvqc_h, stride_kvqc_n, stride_vvq_b,
    stride_vvq_h, stride_vvq_n, stride_vvqi_b, stride_vvqi_h, stride_vvqi_n,
    stride_o_b, stride_o_h, stride_o_m, nheads, seqlen_q, codebook_size,
    CACHE_KEY_SEQLEN_Q, WINDOW_SIZE: tl.constexpr, BLOCK_HEADDIM: tl.
    constexpr, EVEN_M: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.
    constexpr, WRITE_LSE: tl.constexpr):
    start_m = tl.program_id(0)
    off_hb = tl.program_id(1)
    off_b = off_hb // nheads
    off_h = off_hb % nheads
    offs_m = start_m * BLOCK_M + WINDOW_SIZE + BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, BLOCK_HEADDIM)
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) + NEGINF
    acc = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
    Q_block_ptr = tl.make_block_ptr(base=Q + (off_b * stride_q_b + off_h *
        stride_q_h), shape=(seqlen_q, BLOCK_HEADDIM), strides=(stride_q_m, 
        1), offsets=(start_m * BLOCK_M + WINDOW_SIZE + BLOCK_M, 0),
        block_shape=(BLOCK_M, BLOCK_HEADDIM), order=(1, 0))
    K_VQ_block_ptr = tl.make_block_ptr(base=K_VQ + off_h * stride_kvq_h,
        shape=(BLOCK_HEADDIM, codebook_size), strides=(1, stride_kvq_c),
        offsets=(0, 0), block_shape=(BLOCK_HEADDIM, BLOCK_N), order=(0, 1))
    K_VQC_block_ptr = tl.make_block_ptr(base=K_VQ_CNT + (off_b *
        stride_kvqc_b + off_h * stride_kvqc_h + start_m * stride_kvqc_n),
        shape=(codebook_size,), strides=(1,), offsets=(0,), block_shape=(
        BLOCK_N,), order=(0,))
    VVQI_block_ptr = tl.make_block_ptr(base=V_VQ_INDEX + (off_b *
        stride_vvqi_b + off_h * stride_vvqi_h + start_m * stride_vvqi_n),
        shape=(codebook_size,), strides=(1,), offsets=(0,), block_shape=(
        BLOCK_N,), order=(0,))
    V_VQ += off_b * stride_vvq_b + off_h * stride_vvq_h
    if EVEN_M:
        q = tl.load(Q_block_ptr)
    else:
        q = tl.load(Q_block_ptr, boundary_check=(0,), padding_option='zero')
    acc, l_i, m_i = _vq_attn_fwd_inner(acc, l_i, m_i, q, softmax_scale,
        K_VQ_block_ptr, K_VQC_block_ptr, VVQI_block_ptr, V_VQ, stride_vvq_n,
        codebook_size, BLOCK_HEADDIM, BLOCK_N)
    acc = acc / l_i[:, None]
    if WRITE_LSE:
        l_ptrs = L + off_hb * seqlen_q + offs_m
        if EVEN_M:
            tl.store(l_ptrs, m_i + tl.math.log2(l_i))
        else:
            tl.store(l_ptrs, m_i + tl.math.log2(l_i), mask=offs_m < seqlen_q)
    out_ptrs = Out + off_b * stride_o_b + off_h * stride_o_h + (offs_m[:,
        None] * stride_o_m + offs_d[None, :])
    if EVEN_M:
        tl.store(out_ptrs, acc)
    else:
        tl.store(out_ptrs, acc, mask=offs_m[:, None] < seqlen_q)
