@triton.heuristics({'EVEN_M': lambda args: args['seqlen_q'] % args[
    'BLOCK_M'] == 0, 'TOTAL_SLOTS': lambda args: sum(args['CODEBOOK_SIZE'] **
    i for i in range(1, args['CODEBOOK_NUM'] + 1))})
@triton.jit
def _fwd_kernel(Q, CODEBOOK_K, CODEBOOK_V, KCNT, VCNT, Out, softmax_scale,
    stride_qb, stride_qh, stride_qm, stride_kb, stride_kh, stride_kn,
    stride_vb, stride_vh, stride_vn, stride_ob, stride_oh, stride_om,
    nheads, seqlen_q, CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K, CODEBOOK_SIZE:
    tl.constexpr, CODEBOOK_NUM: tl.constexpr, TOTAL_SLOTS: tl.constexpr,
    IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, EVEN_M: tl.
    constexpr, BLOCK_M: tl.constexpr, BLOCK_N: 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 + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    offs_d = tl.arange(0, BLOCK_HEADDIM)
    Q_block_ptr = tl.make_block_ptr(base=Q + (off_b * stride_qb + off_h *
        stride_qh), shape=(seqlen_q, BLOCK_HEADDIM), strides=(stride_qm, 1),
        offsets=(start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_HEADDIM
        ), order=(1, 0))
    K_block_ptr = tl.make_block_ptr(base=K + (off_b * stride_kb + off_h *
        stride_kh), shape=(TOTAL_SLOTS, BLOCK_HEADDIM), strides=(stride_kn,
        1), offsets=(0, 0), block_shape=(BLOCK_N, BLOCK_HEADDIM), order=(1, 0))
    V_block_ptr = tl.make_block_ptr(base=V + (off_b * stride_vb + off_h *
        stride_vh), shape=(TOTAL_SLOTS, BLOCK_HEADDIM), strides=(stride_vn,
        1), offsets=(0, 0), block_shape=(BLOCK_N, BLOCK_HEADDIM), order=(1, 0))
    lse_i = tl.zeros([BLOCK_M], dtype=tl.float32) + NEGINF
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) + NEGINF
    acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
    if EVEN_M:
        q = tl.load(Q_block_ptr)
    else:
        q = tl.load(Q_block_ptr, boundary_check=(0,), padding_option='zero')
    offset = 0
    for level_idx in range(1, CODEBOOK_NUM + 1):
        K_inner_block_ptr = K_block_ptr
        V_inner_block_ptr = V_block_ptr
        z_value = tl.zeros([BLOCK_M], dtype=tl.float32) + NEGINF
        z_idx = tl.zeros([BLOCK_M], dtype=tl.int32)
        for start_n in range(0, CODEBOOK_SIZE, BLOCK_N):
            k = tl.load(K_inner_block_ptr)
            qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
            qk += tl.dot(q, tl.trans(k))
            qk_max = tl.max(qk, 1)
            m_ij = tl.maximum(qk_max * softmax_scale, lse_i)
            p = tl.exp(qk * softmax_scale - m_ij[:, None])
            l_ij = tl.sum(p, 1)
            qk_max_indices = tl.argmax(qk, 1) + start_n
            update_codebook = z_value < qk_max
            z_idx = tl.where(update_codebook, qk_max_indices, z_idx)
            z_value = tl.where()
            acc_o_scale = tl.exp(m_i - m_ij)
            acc_o = acc_o * acc_o_scale[:, None]
            v = tl.load(V_inner_block_ptr)
            p = p.to(v.dtype)
            acc_o += tl.dot(p, v)
            m_i = m_ij
            l_i_new = tl.exp(lse_i - m_ij) + l_ij
            lse_i = m_ij + tl.log(l_i_new)
            K_inner_block_ptr = tl.advance(K_inner_block_ptr, (BLOCK_N, 0))
            V_inner_block_ptr = tl.advance(V_inner_block_ptr, (BLOCK_N, 0))
    end_n = seqlen_k if not IS_CAUSAL else tl.minimum((start_m + 1) *
        BLOCK_M, seqlen_k)
    for start_n in range(0, end_n, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        if EVEN_N & EVEN_M:
            k = tl.load(K_block_ptr)
        else:
            k = tl.load(K_block_ptr, boundary_check=(0,), padding_option='zero'
                )
        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        qk += tl.dot(q, tl.trans(k))
        if not EVEN_N:
            qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0, NEGINF)
        if IS_CAUSAL:
            qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 
                0, NEGINF)
        m_ij = tl.maximum(tl.max(qk, 1) * softmax_scale, lse_i)
        p = tl.exp(qk * softmax_scale - m_ij[:, None])
        l_ij = tl.sum(p, 1)
        acc_o_scale = tl.exp(m_i - m_ij)
        acc_o = acc_o * acc_o_scale[:, None]
        if EVEN_N & EVEN_M:
            v = tl.load(V_block_ptr)
        else:
            v = tl.load(V_block_ptr, boundary_check=(0,), padding_option='zero'
                )
        p = p.to(v.dtype)
        acc_o += tl.dot(p, v)
        m_i = m_ij
        l_i_new = tl.exp(lse_i - m_ij) + l_ij
        lse_i = m_ij + tl.log(l_i_new)
        K_block_ptr = tl.advance(K_block_ptr, (BLOCK_N, 0))
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
    o_scale = tl.exp(m_i - lse_i)
    acc_o = acc_o * o_scale[:, None]
    start_m = tl.program_id(0)
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, BLOCK_HEADDIM)
    out_ptrs = Out + off_b * stride_ob + off_h * stride_oh + (offs_m[:,
        None] * stride_om + offs_d[None, :])
    if EVEN_M:
        tl.store(out_ptrs, acc_o)
    else:
        tl.store(out_ptrs, acc_o, mask=offs_m[:, None] < seqlen_q)
