@triton.jit
def _sparse_attention_decode_fwd_kernel(Q, K, V, sm_scale, Out, L, M,
    stride_qz, stride_qh, stride_qm, stride_qk, stride_kz, stride_kh,
    stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn,
    stride_oz, stride_oh, stride_om, stride_on, Z, H, N_CTX, stride_luth,
    NNZ: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, lut):
    start_n = tl.program_id(0)
    off_hz = tl.program_id(1)
    start_nnz = tl.load(lut + off_hz * stride_luth + start_n)
    kv_offset = off_hz * stride_kh
    q_offset = off_hz * stride_qh
    o_offset = off_hz * stride_oh
    Q_block_ptr = tl.make_block_ptr(base=Q + q_offset, shape=(1,
        BLOCK_DMODEL), strides=(stride_qm, stride_qk), offsets=(0, 0),
        block_shape=(1, BLOCK_DMODEL), order=(1, 0))
    K_block_ptr = tl.make_block_ptr(base=K + kv_offset, shape=(BLOCK_DMODEL,
        N_CTX), strides=(stride_kk, stride_kn), offsets=(0, start_nnz *
        BLOCK_N), block_shape=(BLOCK_DMODEL, BLOCK_N), order=(0, 1))
    V_block_ptr = tl.make_block_ptr(base=V + kv_offset, shape=(N_CTX,
        BLOCK_DMODEL), strides=(stride_vk, stride_vn), offsets=(start_nnz *
        BLOCK_N, 0), block_shape=(BLOCK_N, BLOCK_DMODEL), order=(1, 0))
    O_block_ptr = tl.make_block_ptr(base=Out + o_offset, shape=(NNZ,
        BLOCK_DMODEL), strides=(stride_om, stride_on), offsets=(start_n, 0),
        block_shape=(1, BLOCK_DMODEL), order=(1, 0))
    qk_scale = sm_scale * 1.44269504
    q = tl.load(Q_block_ptr)
    q = (q * qk_scale).to(tl.float16)
    k = tl.load(K_block_ptr)
    qk = tl.expand_dims(tl.sum(tl.trans(q) * k, axis=0), axis=0).to(tl.float32)
    k_indices = start_nnz * BLOCK_N + tl.arange(0, BLOCK_N)
    qk = tl.where(k_indices < N_CTX, qk, float('-inf'))
    m = tl.max(qk, axis=1, return_indices=False)
    p = tl.math.exp2(qk - m)
    l = tl.sum(p, axis=1)
    v = tl.load(V_block_ptr).to(tl.float16)
    p = p.to(tl.float16)
    acc = tl.expand_dims(tl.sum(tl.trans(p) * v, axis=0), axis=0)
    tl.store(O_block_ptr, acc)
    l_ptrs = L + off_hz * NNZ + start_n + tl.arange(0, 1)
    m_ptrs = M + off_hz * NNZ + start_n + tl.arange(0, 1)
    tl.store(l_ptrs, l)
    tl.store(m_ptrs, m)
