chore: import upstream snapshot with attribution

ReSA/llm/kernel/flash_attention_with_kv_cache.py

@@ -0,0 +1,351 @@
+import math
+import torch
+
+import triton
+import triton.language as tl
+
+def is_hip():
+    return triton.runtime.driver.active.get_current_target().backend == "hip"
+
+
+def num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, num_m_blocks, size_one_kv_head,
+                         is_causal_or_local, max_splits):
+    """
+    Determines the optimal number of splits for maximizing GPU occupancy while balancing memory efficiency.
+
+    Parameters:
+    - total_mblocks (int): Total number of m_blocks.
+    - num_SMs (int): Number of Streaming Multiprocessors (SMs) in the GPU.
+    - num_n_blocks (int): Number of n_blocks.
+    - num_m_blocks (int): Number of m_blocks.
+    - size_one_kv_head (int): Size of one KV head in bytes.
+    - is_causal_or_local (bool): Indicates whether the operation is causal or local.
+    - max_splits (int): Maximum number of allowed splits.
+
+    Returns:
+    - int: The optimal number of splits.
+    """
+    # If we have enough m_blocks to almost fill the SMs, prefer 1 split unless memory constraints apply.
+    if total_mblocks >= 0.8 * num_SMs:
+        size_l2 = 50 * 1024 * 1024  # L2 cache size assumption (50MB)
+        # Only split if each KV head is too large for L2 and there are enough m_blocks
+        if size_one_kv_head > size_l2 and num_m_blocks >= num_SMs * 2 and not is_causal_or_local:
+            return min((size_one_kv_head + size_l2 - 1) // size_l2, max_splits)
+        else:
+            return 1
+
+    # If num_n_blocks is too small, we don't split
+    if num_n_blocks <= 4:
+        return 1
+
+    # Limit max_splits to a reasonable range
+    max_splits = min(max_splits, num_SMs, num_n_blocks)
+
+    max_efficiency = 0.0
+    efficiency = []
+
+    # Compute efficiency for different splits
+    for num_splits in range(1, max_splits + 1):
+        n_waves = (total_mblocks * num_splits) / num_SMs
+        eff = n_waves / math.ceil(n_waves)
+        # Track max efficiency
+        if eff > max_efficiency:
+            max_efficiency = eff
+
+        efficiency.append(eff)
+
+    # Find the smallest number of splits that achieves at least 85% of max efficiency
+    for num_splits in range(1, max_splits + 1):
+        if efficiency[num_splits - 1] >= 0.95 * max_efficiency:
+            return num_splits
+
+    return 1
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [1, 2, 4, 8, 16]
+    ],
+    key=['gqa_group_size', 'BLOCK_H', 'BLOCK_N', 'BLOCK_D', 'BLOCK_V'],
+)
+@triton.jit
+def _fwd_kernel_with_kv_cache(
+    Q, K, V, Out, L,
+    sm_scale,
+    cache_seqlens,
+    stride_qz, stride_qt, stride_qh, stride_qd,
+    stride_kz, stride_kt, stride_kh, stride_kd,
+    stride_vz, stride_vt, stride_vh, stride_vd,
+    stride_oz, stride_ot, stride_oh, stride_os, stride_od,
+    stride_lz, stride_lt, stride_lh, stride_ls,
+    num_splits: tl.constexpr,
+    seqlen_q: tl.constexpr,
+    num_m_blocks: tl.constexpr,
+    gqa_group_size: tl.constexpr,
+    BLOCK_H: tl.constexpr, 
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+):
+    off_sm = tl.program_id(0).to(tl.int64)
+    off_split, off_m = off_sm // num_m_blocks, off_sm % num_m_blocks
+    off_h_for_kv = tl.program_id(1).to(tl.int64)
+    off_z = tl.program_id(2).to(tl.int64)
+
+    off_h_q = off_h_for_kv * gqa_group_size
+    offs_h = tl.arange(0, BLOCK_H)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_D)
+    offs_v = tl.arange(0, BLOCK_V)
+    mask_h = offs_h < gqa_group_size
+    seqlen_k = tl.load(cache_seqlens + off_z)
+
+    Q += off_z * stride_qz + off_h_q * stride_qh
+    K += off_z * stride_kz + off_h_for_kv * stride_kh
+    V += off_z * stride_vz + off_h_for_kv * stride_vh
+    L += off_z * stride_lz + off_h_q * stride_lh + off_split * stride_ls
+    Out += off_z * stride_oz + off_h_q * stride_oh + off_split * stride_os
+    
+    num_kv_blocks = tl.cdiv(seqlen_k, BLOCK_N)
+    blocks_per_split = num_kv_blocks // num_splits
+    remaining_blocks = num_kv_blocks % num_splits
+    loop_range = blocks_per_split + (1 if off_split < remaining_blocks else 0)
+    start = blocks_per_split * off_split + min(off_split, remaining_blocks)
+
+    offs_m = tl.arange(0, BLOCK_M) + off_m * BLOCK_M
+    mask_q = offs_m < seqlen_q
+    mask_qh = mask_q[:, None] & mask_h[None, :]
+    q_idx = offs_m[:, None] + tl.zeros([BLOCK_H], dtype=tl.int32) + seqlen_k - seqlen_q
+    q_idx = tl.reshape(q_idx, [BLOCK_M * BLOCK_H])
+
+    q = tl.load(Q + offs_m[:, None, None] * stride_qt + offs_h[None, :, None] * stride_qh + offs_d[None, None, :] * stride_qd, mask=mask_qh[:, :, None]) ## padding to min 16
+    q = tl.reshape(q, (BLOCK_M * BLOCK_H, BLOCK_D))
+
+    m_i = tl.full([BLOCK_M * BLOCK_H], float("-inf"), dtype=tl.float32)
+    l_i = tl.full([BLOCK_M * BLOCK_H], 1.0, dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M * BLOCK_H, BLOCK_V], dtype=tl.float32)
+
+    k_ptrs = K + offs_n[None, :] * stride_kt + offs_d[:, None] * stride_kd
+    v_ptrs = V + offs_n[:, None] * stride_vt + offs_v[None, :] * stride_vd
+
+    for block_idx in range(start, start + loop_range):
+        start_n = block_idx * BLOCK_N
+        k = tl.load(k_ptrs + start_n * stride_kt, mask=offs_n[None, :] + start_n < seqlen_k, cache_modifier=".ca")
+        qk = tl.dot(q, k)
+
+        causal_mask = q_idx[:, None] >= start_n + offs_n[None, :]
+        qk = tl.where(causal_mask, qk, -1.0e6)
+        qk *= sm_scale
+
+        m_ij = tl.maximum(m_i, tl.max(qk, 1))
+        qk -= m_ij[:, None]
+        p = tl.exp(qk)
+        l_ij = tl.sum(p, 1)
+        alpha = tl.exp(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        acc = acc * alpha[:, None]
+        
+        v = tl.load(v_ptrs + start_n * stride_vt, mask=offs_n[:, None] + start_n < seqlen_k, cache_modifier=".ca")
+        p = p.to(v.type.element_ty)
+
+        acc += tl.dot(p, v)
+        m_i = m_ij
+
+    l_recip = 1 / l_i[:, None]
+    acc = acc * l_recip
+    m_i += tl.math.log(l_i)
+
+    l_ptrs = L + offs_m[:, None] * stride_lt + offs_h * stride_lh
+    m_i = tl.reshape(m_i, (BLOCK_M, BLOCK_H))
+    tl.store(l_ptrs, m_i, mask=mask_qh)
+    O_ptrs = Out + offs_m[:, None, None] * stride_ot + offs_h[None, :, None] * stride_oh + offs_v[None, None, :] * stride_od
+    acc = tl.reshape(acc, (BLOCK_M, BLOCK_H, BLOCK_V))
+    tl.store(O_ptrs, acc, mask=mask_qh[:, :, None])
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [1, 2, 4, 8, 16]
+    ],
+    key=['BLOCK_V'],
+)
+@triton.jit
+def combine(
+    out_partial, out, L,
+    stride_op_z, stride_op_t, stride_op_h, stride_op_s, stride_op_d,
+    stride_o_z, stride_o_t, stride_o_h, stride_o_d,
+    stride_l_z, stride_l_t, stride_l_h, stride_l_s,
+    num_splits: tl.constexpr,
+    num_splits_pow2: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+):
+    off_h = tl.program_id(0).to(tl.int64)
+    off_t = tl.program_id(1).to(tl.int64)
+    off_z = tl.program_id(2).to(tl.int64)
+
+    split = tl.arange(0, num_splits_pow2)
+    split_mask = split < num_splits
+
+    L += off_z * stride_l_z + off_t * stride_l_t + off_h * stride_l_h
+    out_partial += off_z * stride_op_z + off_t * stride_op_t + off_h * stride_op_h
+    out += off_z * stride_o_z + off_t * stride_o_t + off_h * stride_o_h
+
+    lse_local = tl.load(L + split * stride_l_s, mask=split_mask, other=float("-inf"))
+    lse_max_local = tl.max(lse_local, axis=0)
+
+    lse_logsum_local = tl.sum(tl.exp(lse_local - lse_max_local), axis=0)
+    lse_logsum_local = tl.log(lse_logsum_local) + lse_max_local
+    
+    po_local = tl.load(out_partial + split[:, None] * stride_op_s + tl.arange(0, BLOCK_V) * stride_op_d, mask=split_mask[:, None])
+    scale_local = tl.exp(lse_local - lse_logsum_local)
+    accum_local = tl.sum(po_local * scale_local[:, None], axis=0)
+    tl.store(out + tl.arange(0, BLOCK_V) * stride_o_d, accum_local)
+
+def flash_attention_with_kv_cache(
+    q, k, v,
+    cache_seqlens,
+    sm_scale=None,
+):
+    # split q to blocks
+    batch, seqlen_q, n_heads, key_dim = q.shape
+    _, _, n_kv_heads, head_dim = v.shape
+    gqa_group_size = n_heads // n_kv_heads
+    block_h = triton.next_power_of_2(gqa_group_size)
+    block_m = max(256 // block_h, 1)
+    block_n = 32
+
+    # assert seqlen_q <= 32, "it seems the performance is not good when seqlen_q > 32"
+
+    assert k.size(0) == v.size(0)
+    assert q.size(3) == k.size(3)
+    assert k.size(1) == v.size(1)
+    assert key_dim in {64, 128, 256}
+    assert head_dim in {64, 128, 256}
+
+    props = torch.cuda.get_device_properties(torch.device("cuda:0"))
+    num_sm = props.multi_processor_count
+    num_m_blocks = triton.cdiv(seqlen_q, block_m)
+    num_n_blocks = triton.cdiv(cache_seqlens.max(), block_n)
+    size_one_kv_head = cache_seqlens.max() * block_n * (key_dim + head_dim) * 2
+    total_mblocks = batch * n_kv_heads
+    num_splits = num_splits_heuristic(
+        total_mblocks, num_sm, num_n_blocks, num_m_blocks,
+        size_one_kv_head, is_causal_or_local=True, max_splits=16
+    )
+
+    out_partial = torch.empty((batch, seqlen_q, n_heads, num_splits, head_dim), device=q.device, dtype=torch.float32)
+    out = torch.empty((batch, seqlen_q, n_heads, head_dim), device=q.device, dtype=q.dtype)
+    L = torch.empty((batch, seqlen_q, n_heads, num_splits), device=q.device, dtype=torch.float32)
+
+    if is_hip():
+        extra_kern_args = {"waves_per_eu": 1}
+    else:
+        extra_kern_args = {}
+
+    with torch.cuda.device(q.device.index): 
+        grid = lambda META: (num_splits * num_m_blocks, n_kv_heads, batch)
+        _fwd_kernel_with_kv_cache[grid](
+            q, k, v, out_partial, L,
+            sm_scale if sm_scale is not None else key_dim ** -0.5,
+            cache_seqlens.contiguous(),
+            *q.stride(),
+            *k.stride(),
+            *v.stride(),
+            *out_partial.stride(),
+            *L.stride(),
+            num_splits=num_splits,
+            seqlen_q=seqlen_q,
+            num_m_blocks=num_m_blocks,
+            gqa_group_size=gqa_group_size,
+            BLOCK_H = block_h,
+            BLOCK_M = block_m,
+            BLOCK_N = block_n,
+            BLOCK_D = key_dim,
+            BLOCK_V = head_dim,
+            **extra_kern_args
+        )
+        grid = lambda META: (n_heads, seqlen_q, batch)
+        combine[grid](
+            out_partial, out, L,
+            *out_partial.stride(),
+            *out.stride(),
+            *L.stride(),
+            num_splits=num_splits,
+            num_splits_pow2=triton.next_power_of_2(num_splits),
+            BLOCK_V=head_dim,
+            **extra_kern_args
+        )
+    return out
+
+def ref_program_fa(query, key, value, cache_seqlens):
+    # latency reference
+    # from flash_attn_interface import flash_attn_with_kvcache, flash_attn_func # fa3
+    from flash_attn import flash_attn_with_kvcache, flash_attn_func #fa2
+    output = flash_attn_with_kvcache(query, key, value, cache_seqlens=cache_seqlens, causal=True)
+    return output
+
+
+
+def debug(name,expect, actual, atol=1e-3, rtol=1e-3):
+    all_close = torch.allclose(expect, actual, atol=atol, rtol=rtol)
+    print(name + "  all_close={}".format(all_close))
+    if not all_close:
+        # print(expect[3, 28])
+        # print(actual[3, 28])
+        diff = (expect - actual).abs()
+        print("all_close={}, max={}, min={}, mean={}".format(all_close, diff.max().item(), diff.min().item(), diff.mean().item()))
+        max_indices = torch.nonzero(diff == diff.max().item())
+        first_index = tuple(max_indices[0].tolist())
+        print(f"Index: {first_index}, expect: {expect[first_index]}, actual: {actual[first_index]}") 
+
+if __name__ == "__main__":
+    import argparse
+    import time
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=8, help='batch size')
+    parser.add_argument('--seqlen_q', type=int, default=128, help='sequence length')
+    parser.add_argument('--heads', type=int, default=28, help='heads')
+    parser.add_argument('--heads_kv', type=int, default=4, help='heads_kv')
+    parser.add_argument('--max_cache_seqlen', type=int, default=65536, help='kvcache sequence length')
+    parser.add_argument('--dim', type=int, default=128, help='dim')
+    parser.add_argument('--dim_v', type=int, default=128, help='dim_v')
+    parser.add_argument('--load_from_file', type=str, default=None, help='load from file')
+    args = parser.parse_args()
+
+    batch, seqlen_q, heads, heads_kv, max_cache_seqlen, dim, dim_v = args.batch, args.seqlen_q, args.heads, args.heads_kv, args.max_cache_seqlen, args.dim, args.dim_v
+
+    dtype = torch.bfloat16
+
+    Q = torch.randn((batch, seqlen_q, heads, dim), dtype=dtype, device='cuda')
+    K = torch.randn((batch, max_cache_seqlen, heads_kv, dim), dtype=dtype, device='cuda')
+    V = torch.randn((batch, max_cache_seqlen, heads_kv, dim_v), dtype=dtype, device='cuda')
+
+    cache_seqlens = torch.randint(max_cache_seqlen - 32, max_cache_seqlen, (batch,), device='cuda', dtype=torch.int32)
+    print("cache_seqlens: ", cache_seqlens)
+    # parity reference
+    ref = ref_program_fa(Q, K, V, cache_seqlens)
+    # ref = ref_program_triton(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, max_num_blocks, block_size)
+    # out = kernel(Q, K, V, block_indices, cache_seqlens, actual_num_blocks, glse, Output_partial)
+    # out = sparse_gqa_decode_varlen_indice(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, block_size)
+    out = flash_attention_with_kv_cache(Q, K, V, cache_seqlens)
+    debug("output", ref, out, atol=1e-3, rtol=1e-3)
+
+    ## latency reference
+    for i in range(10):
+        ref = ref_program_fa(Q, K, V, cache_seqlens)
+    torch.cuda.synchronize()
+    start = time.time()
+    for i in range(100):
+        ref = ref_program_fa(Q, K, V, cache_seqlens)
+    torch.cuda.synchronize()
+    print("dense time: ", (time.time() - start) / 100*1000)
+
+    for i in range(10):
+        out = flash_attention_with_kv_cache(Q, K, V, cache_seqlens)
+    torch.cuda.synchronize()
+    start = time.time()
+    for i in range(100):
+        out = flash_attention_with_kv_cache(Q, K, V, cache_seqlens)
+    torch.cuda.synchronize()
+    print("sparse time: ", (time.time() - start) / 100*1000)

ReSA/llm/kernel/flash_sparse_decoding.py

@@ -0,0 +1,315 @@
+import math
+import torch
+
+import triton
+import triton.language as tl
+
+def is_hip():
+    return triton.runtime.driver.active.get_current_target().backend == "hip"
+
+
+def num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, num_m_blocks, size_one_kv_head,
+                         is_causal_or_local, max_splits):
+    """
+    Determines the optimal number of splits for maximizing GPU occupancy while balancing memory efficiency.
+
+    Parameters:
+    - total_mblocks (int): Total number of m_blocks.
+    - num_SMs (int): Number of Streaming Multiprocessors (SMs) in the GPU.
+    - num_n_blocks (int): Number of n_blocks.
+    - num_m_blocks (int): Number of m_blocks.
+    - size_one_kv_head (int): Size of one KV head in bytes.
+    - is_causal_or_local (bool): Indicates whether the operation is causal or local.
+    - max_splits (int): Maximum number of allowed splits.
+
+    Returns:
+    - int: The optimal number of splits.
+    """
+    # If we have enough m_blocks to almost fill the SMs, prefer 1 split unless memory constraints apply.
+    if total_mblocks >= 0.8 * num_SMs:
+        size_l2 = 50 * 1024 * 1024  # L2 cache size assumption (50MB)
+        # Only split if each KV head is too large for L2 and there are enough m_blocks
+        if size_one_kv_head > size_l2 and num_m_blocks >= num_SMs * 2 and not is_causal_or_local:
+            return min((size_one_kv_head + size_l2 - 1) // size_l2, max_splits)
+        else:
+            return 1
+
+    # If num_n_blocks is too small, we don't split
+    if num_n_blocks <= 4:
+        return 1
+
+    # Limit max_splits to a reasonable range
+    max_splits = min(max_splits, num_SMs, num_n_blocks)
+
+    max_efficiency = 0.0
+    efficiency = []
+
+    # Compute efficiency for different splits
+    for num_splits in range(1, max_splits + 1):
+        n_waves = (total_mblocks * num_splits) / num_SMs
+        eff = n_waves / math.ceil(n_waves)
+        # Track max efficiency
+        if eff > max_efficiency:
+            max_efficiency = eff
+
+        efficiency.append(eff)
+
+    # Find the smallest number of splits that achieves at least 85% of max efficiency
+    for num_splits in range(1, max_splits + 1):
+        if efficiency[num_splits - 1] >= 0.95 * max_efficiency:
+            return num_splits
+
+    return 1
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [1, 2, 4, 8, 16]
+    ],
+    key=['gqa_group_size', 'BLOCK_H', 'BLOCK_N', 'BLOCK_D', 'BLOCK_V'],
+)
+@triton.jit
+def _fwd_kernel_decoding(
+    Q, K, V, Out, L,
+    sm_scale,
+    cache_seqlens,
+    block_indices_ptr,
+    stride_qz, stride_qh, stride_qd,
+    stride_kz, stride_kt, stride_kh, stride_kd,
+    stride_vz, stride_vt, stride_vh, stride_vd,
+    stride_oz, stride_oh, stride_os, stride_od,
+    stride_lz, stride_lh, stride_ls,
+    stride_bz, stride_bn, stride_bd,
+    max_selected_blocks: tl.constexpr,
+    num_splits: tl.constexpr,
+    gqa_group_size: tl.constexpr,
+    BLOCK_H: tl.constexpr, 
+    BLOCK_N: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+):
+    off_z = tl.program_id(0).to(tl.int64)
+    off_h_for_kv = tl.program_id(1).to(tl.int64)
+    off_split = tl.program_id(2).to(tl.int64)
+
+    off_h_q = off_h_for_kv * gqa_group_size
+    
+    offs_m = tl.arange(0, BLOCK_H) ## head 
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_D)
+    offs_v = tl.arange(0, BLOCK_V)
+
+    seqlen_k = tl.load(cache_seqlens + off_z)
+
+    Q += off_z * stride_qz + off_h_q * stride_qh
+    K += off_z * stride_kz + off_h_for_kv * stride_kh
+    V += off_z * stride_vz + off_h_for_kv * stride_vh
+    L += off_z * stride_lz + off_h_q * stride_lh + off_split * stride_ls
+    Out += off_z * stride_oz + off_h_q * stride_oh + off_split * stride_os
+    block_indices_ptr += off_z * stride_bz + off_h_for_kv * stride_bn
+
+    q = tl.load(Q + offs_m[:, None] * stride_qh + offs_d[None, :] * stride_qd,
+                mask=(offs_m[:, None] < gqa_group_size)) ## padding to min 16
+
+    blocks_per_split = max_selected_blocks // num_splits
+    remaining_blocks = max_selected_blocks % num_splits
+    loop_range = blocks_per_split + (1 if off_split < remaining_blocks else 0)
+    start = blocks_per_split * off_split + min(off_split, remaining_blocks)
+    m_i = tl.full([BLOCK_H], float("-inf"), dtype=tl.float32)
+    l_i = tl.full([BLOCK_H], 1.0, dtype=tl.float32)
+    acc = tl.zeros([BLOCK_H, BLOCK_V], dtype=tl.float32)
+
+    k_ptrs = K + offs_n[None, :] * stride_kt + offs_d[:, None] * stride_kd
+    v_ptrs = V + offs_n[:, None] * stride_vt + offs_v[None, :] * stride_vd
+    for block_ptr_idx in range(start, start + loop_range):
+        block_idx = tl.load(block_indices_ptr + block_ptr_idx * stride_bd)
+        if block_idx >= 0:
+            start_n = block_idx * BLOCK_N
+            k = tl.load(k_ptrs + start_n * stride_kt, mask=offs_n[None, :] + start_n < seqlen_k)
+            qk = tl.dot(q, k)
+            qk = tl.where(offs_n[None, :] + start_n < seqlen_k, qk, -1e6)
+            qk *= sm_scale
+
+            m_ij = tl.maximum(m_i, tl.max(qk, 1))
+            qk -= m_ij[:, None]
+            p = tl.exp(qk)
+            l_ij = tl.sum(p, 1)
+            alpha = tl.exp(m_i - m_ij)
+            l_i = l_i * alpha + l_ij
+            acc = acc * alpha[:, None]
+            
+            v = tl.load(v_ptrs + start_n * stride_vt, mask=offs_n[:, None] + start_n < seqlen_k)
+            p = p.to(v.type.element_ty)
+
+            acc += tl.dot(p, v)
+            m_i = m_ij
+
+    l_recip = 1 / l_i[:, None]
+    acc = acc * l_recip
+    m_i += tl.math.log(l_i)
+
+    l_ptrs = L + offs_m * stride_lh
+    tl.store(l_ptrs, m_i, mask=(offs_m < gqa_group_size))
+    O_ptrs = Out + offs_m[:, None] * stride_oh + offs_v[None, :] * stride_od
+    tl.store(O_ptrs, acc, mask=(offs_m[:, None] < gqa_group_size))
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [1, 2, 4, 8, 16]
+    ],
+    key=['BLOCK_V'],
+)
+@triton.jit
+def combine(
+    out_partial, out, L,
+    stride_op_z, stride_op_h, stride_op_s, stride_op_d,
+    stride_o_z, stride_o_h, stride_o_d,
+    stride_l_z, stride_l_h, stride_l_s,
+    num_splits: tl.constexpr,
+    num_splits_pow2: tl.constexpr,
+    BLOCK_V: tl.constexpr,
+):
+    off_z = tl.program_id(0).to(tl.int64)
+    off_h = tl.program_id(1).to(tl.int64)
+
+    split = tl.arange(0, num_splits_pow2)
+    split_mask = split < num_splits
+
+    lse_local = tl.load(L + off_z * stride_l_z + off_h * stride_l_h + split * stride_l_s, mask=split_mask, other=float("-inf"))
+    lse_max_local = tl.max(lse_local, axis=0)
+
+    lse_logsum_local = tl.sum(tl.exp(lse_local - lse_max_local), axis=0)
+    lse_logsum_local = tl.log(lse_logsum_local) + lse_max_local
+    
+    po_local = tl.load(out_partial + off_z * stride_op_z + off_h * stride_op_h + split[:, None] * stride_op_s + tl.arange(0, BLOCK_V) * stride_op_d, mask=split_mask[:, None])
+    scale_local = tl.exp(lse_local - lse_logsum_local)
+    accum_local = tl.sum(po_local * scale_local[:, None], axis=0)
+    tl.store(out + off_z * stride_o_z + off_h * stride_o_h + tl.arange(0, BLOCK_V) * stride_o_d, accum_local)
+
+def flash_block_sparse_decoding(
+    q, k, v,
+    cache_seqlens,
+    block_indices,
+    sm_scale=None,
+    block_size=64,
+    num_splits=None
+):
+    # split q to blocks
+    batch, n_heads, key_dim = q.shape
+    _, _, n_kv_heads, head_dim = v.shape
+    gqa_group_size = n_heads // n_kv_heads
+    max_selected_blocks = block_indices.shape[-1]
+    block_h = max(triton.next_power_of_2(gqa_group_size), 16)
+
+    assert k.size(0) == v.size(0)
+    assert q.size(2) == k.size(3)
+    assert k.size(1) == v.size(1)
+    assert key_dim in {64, 128, 256}
+    assert head_dim in {64, 128, 256}
+    assert triton.next_power_of_2(block_size) == block_size, "block size must be power of 2"
+
+    props = torch.cuda.get_device_properties(torch.device("cuda:0"))
+    num_sm = props.multi_processor_count
+    num_m_blocks = 1
+    num_n_blocks = max_selected_blocks
+    size_one_kv_head = max_selected_blocks * block_size * (key_dim + head_dim) * 2
+    total_mblocks = batch * n_kv_heads * num_m_blocks
+    if num_splits is None:
+        num_splits = num_splits_heuristic(
+            total_mblocks, num_sm, num_n_blocks, num_m_blocks,
+            size_one_kv_head, is_causal_or_local=True, max_splits=8)
+
+    out_partial = torch.empty((batch, n_heads, num_splits, head_dim), device=q.device, dtype=torch.float32)
+    out = torch.empty((batch, n_heads, head_dim), device=q.device, dtype=q.dtype)
+    L = torch.empty((batch, n_heads, num_splits), device=q.device, dtype=torch.float32)
+
+    if is_hip():
+        extra_kern_args = {"waves_per_eu": 1}
+    else:
+        extra_kern_args = {}
+
+    with torch.cuda.device(q.device.index): 
+        grid = lambda META: (batch, n_kv_heads, num_splits)
+        _fwd_kernel_decoding[grid](
+            q, k, v, out_partial, L,
+            sm_scale if sm_scale is not None else key_dim ** -0.5,
+            cache_seqlens.contiguous(),
+            block_indices.contiguous(),
+            *q.stride(),
+            *k.stride(),
+            *v.stride(),
+            *out_partial.stride(),
+            *L.stride(),
+            *block_indices.stride(),
+            max_selected_blocks=max_selected_blocks,
+            num_splits=num_splits,
+            gqa_group_size=gqa_group_size,
+            BLOCK_H = block_h,
+            BLOCK_N = block_size,
+            BLOCK_D = key_dim,
+            BLOCK_V = head_dim,
+            **extra_kern_args
+        )
+        grid = lambda META: (batch, n_heads)
+        combine[grid](
+            out_partial, out, L,
+            *out_partial.stride(),
+            *out.stride(),
+            *L.stride(),
+            num_splits=num_splits,
+            num_splits_pow2=triton.next_power_of_2(num_splits),
+            BLOCK_V = head_dim,
+            **extra_kern_args
+        )
+    return out
+
+def main():
+    from torch.nn import functional as F
+    import time
+    torch.cuda.manual_seed(0)
+    bsz, n_head, key_dim = 4, 2, 128
+    n_kv_seq = 8192
+    head_dim = 128
+    gqa_size = 6
+    block_size = 16
+    dtype = torch.float16
+    xq = torch.randn((bsz, n_head * gqa_size, key_dim), device='cuda', dtype=dtype)
+    xk = torch.randn((bsz, n_kv_seq, n_head, key_dim), device='cuda', dtype=dtype)
+    xv = torch.randn((bsz, n_kv_seq, n_head, head_dim), device='cuda', dtype=dtype)
+    cache_seqlens = torch.randint(100, n_kv_seq, (bsz,), device='cuda', dtype=torch.int32)
+    sparse_mask = torch.rand((bsz, n_head, (n_kv_seq + block_size - 1) // block_size), device='cuda') > 0.9
+    max_selected_blocks = sparse_mask.sum(dim=-1).max()
+    print("max_selected_blocks", max_selected_blocks)
+    sparse_indices = torch.full((bsz, n_head, max_selected_blocks), -1, device='cuda', dtype=torch.int32)
+    for i in range(bsz):
+        for j in range(n_head):
+            valid_blocks = torch.where(sparse_mask[i, j])[0]
+            sparse_indices[i, j, :len(valid_blocks)] = valid_blocks
+    torch.cuda.synchronize()
+    start_time = time.time()
+    for _ in range(100):
+        triton_output = flash_block_sparse_decoding(xq, xk, xv, cache_seqlens, sparse_indices, block_size=block_size)
+    torch.cuda.synchronize()
+    end_time = time.time()
+    print(f"Triton Time taken: {end_time - start_time} seconds")
+    
+    naive_mask = torch.zeros((bsz, n_head, 1, n_kv_seq), device=xq.device, dtype=torch.bool)
+    for i in range(bsz):
+        block_mask = sparse_mask[i].repeat_interleave(block_size, dim=-1)
+        block_mask = torch.masked_fill(block_mask, torch.arange(n_kv_seq, device=xq.device) >= cache_seqlens[i], False)
+        naive_mask[i] = block_mask.unsqueeze(1)
+    
+    torch.cuda.synchronize()
+    start_time = time.time()
+    for _ in range(100):
+        output = F.scaled_dot_product_attention(xq.unsqueeze(2), xk.transpose(1, 2), xv.transpose(1, 2), attn_mask=naive_mask.repeat_interleave(gqa_size, dim=1), enable_gqa=True)
+        output = output.view(bsz, n_head * gqa_size, head_dim)
+    torch.cuda.synchronize()
+    end_time = time.time()
+    print(f"Torch SDPA Time taken: {end_time - start_time} seconds")
+    print(output.shape, triton_output.shape)
+    print((output - triton_output).abs().max(), (output - triton_output).abs().mean())
+    
+if __name__ == "__main__":
+    main()

ReSA/llm/kernel/rotary.py

@@ -0,0 +1,354 @@
+# Copyright (c) 2023, Tri Dao.
+
+from typing import Optional, Union
+
+import torch
+
+import triton
+import triton.language as tl
+
+from typing import Optional, Union
+from einops import rearrange, repeat
+
+
+def rotate_half(x, interleaved=False):
+    if not interleaved:
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+    else:
+        x1, x2 = x[..., ::2], x[..., 1::2]
+        return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2)
+
+
+def apply_rotary_emb_torch(x, cos, sin, interleaved=False, inplace=False):
+    """
+    x: (batch_size, seqlen, nheads, headdim)
+    cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
+    """
+    ro_dim = cos.shape[-1] * 2
+    assert ro_dim <= x.shape[-1]
+    cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
+    sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
+    return torch.cat(
+        [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]],
+        dim=-1,
+    )
+
+@triton.jit
+def rotary_kernel(
+    OUT,  # Pointers to matrices
+    X,
+    COS,
+    SIN,
+    CU_SEQLENS,
+    SEQLEN_OFFSETS,  # this could be int or a pointer
+    # Matrix dimensions
+    seqlen,
+    rotary_dim,
+    seqlen_ro,
+    # strides
+    stride_out_batch,
+    stride_out_seqlen,
+    stride_out_nheads,
+    stride_out_headdim,
+    stride_x_batch,
+    stride_x_seqlen,
+    stride_x_nheads,
+    stride_x_headdim,
+    # Meta-parameters
+    BLOCK_K: tl.constexpr,
+    IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    INTERLEAVED: tl.constexpr,
+    CONJUGATE: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+):
+    pid_m = tl.program_id(axis=0)
+    pid_batch = tl.program_id(axis=1)
+    pid_head = tl.program_id(axis=2)
+    rotary_dim_half = rotary_dim // 2
+
+    if not IS_VARLEN:
+        X = X + pid_batch * stride_x_batch + pid_head * stride_x_nheads
+        OUT = OUT + pid_batch * stride_out_batch + pid_head * stride_out_nheads
+    else:
+        start_idx = tl.load(CU_SEQLENS + pid_batch)
+        seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx
+        X = X + start_idx * stride_x_seqlen + pid_head * stride_x_nheads
+        OUT = OUT + start_idx * stride_out_seqlen + pid_head * stride_out_nheads
+
+    if pid_m * BLOCK_M >= seqlen:
+        return
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    if not IS_SEQLEN_OFFSETS_TENSOR:
+        rm_cs = rm + SEQLEN_OFFSETS
+    else:
+        rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch)
+    rk = tl.arange(0, BLOCK_K)
+    rk_half = tl.arange(0, BLOCK_K // 2)
+
+    if not INTERLEAVED:
+        # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT
+        X = X + (rm[:, None] * stride_x_seqlen + rk_half[None, :] * stride_x_headdim)
+        COS = COS + (rm_cs[:, None] * rotary_dim_half + rk_half[None, :])
+        SIN = SIN + (rm_cs[:, None] * rotary_dim_half + rk_half[None, :])
+        cos = tl.load(
+            COS, mask=(rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < rotary_dim_half), other=1.0
+        ).to(tl.float32)
+        sin = tl.load(
+            SIN, mask=(rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < rotary_dim_half), other=0.0
+        ).to(tl.float32)
+        x0 = tl.load(
+            X, mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half), other=0.0
+        ).to(tl.float32)
+        x1 = tl.load(
+            X + rotary_dim_half * stride_x_headdim,
+            mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half),
+            other=0.0,
+        ).to(tl.float32)
+        if CONJUGATE:
+            sin = -sin
+        o0 = x0 * cos - x1 * sin
+        o1 = x0 * sin + x1 * cos
+        # write back result
+        OUT = OUT + (rm[:, None] * stride_out_seqlen + rk_half[None, :] * stride_out_headdim)
+        tl.store(OUT, o0, mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half))
+        tl.store(
+            OUT + rotary_dim_half * stride_out_headdim,
+            o1,
+            mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half),
+        )
+    else:
+        # We don't want to load X[0, 2, 4, ...] and X[1, 3, 5, ...] separately since both are slow.
+        # Instead, we load x0 = X[0, 1, 2, 3, ...] and x1 = X[1, 0, 3, 2, ...].
+        # Loading x0 will be fast but x1 will be slow.
+        # Then we load cos = COS[0, 0, 1, 1, ...] and sin = SIN[0, 0, 1, 1, ...].
+        # Then we do the calculation and use tl.where to pick put the right outputs for the even
+        # and for the odd indices.
+        rk_swap = rk + ((rk + 1) % 2) * 2 - 1  # 1, 0, 3, 2, 5, 4, ...
+        rk_repeat = tl.arange(0, BLOCK_K) // 2
+        X0 = X + (rm[:, None] * stride_x_seqlen + rk[None, :] * stride_x_headdim)
+        X1 = X + (rm[:, None] * stride_x_seqlen + rk_swap[None, :] * stride_x_headdim)
+        COS = COS + (rm_cs[:, None] * rotary_dim_half + rk_repeat[None, :])
+        SIN = SIN + (rm_cs[:, None] * rotary_dim_half + rk_repeat[None, :])
+        cos = tl.load(
+            COS,
+            mask=(rm_cs[:, None] < seqlen_ro) & (rk_repeat[None, :] < rotary_dim_half),
+            other=1.0,
+        ).to(tl.float32)
+        sin = tl.load(
+            SIN,
+            mask=(rm_cs[:, None] < seqlen_ro) & (rk_repeat[None, :] < rotary_dim_half),
+            other=0.0,
+        ).to(tl.float32)
+        x0 = tl.load(X0, mask=(rm[:, None] < seqlen) & (rk[None, :] < rotary_dim), other=0.0).to(
+            tl.float32
+        )
+        x1 = tl.load(
+            X1, mask=(rm[:, None] < seqlen) & (rk_swap[None, :] < rotary_dim), other=0.0
+        ).to(tl.float32)
+        if CONJUGATE:
+            sin = -sin
+        x0_cos = x0 * cos
+        x1_sin = x1 * sin
+        out = tl.where(rk[None, :] % 2 == 0, x0_cos - x1_sin, x0_cos + x1_sin)
+        OUT = OUT + (rm[:, None] * stride_out_seqlen + rk[None, :] * stride_out_headdim)
+        tl.store(OUT, out, mask=(rm[:, None] < seqlen) & (rk[None, :] < rotary_dim))
+
+
+def apply_rotary(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    seqlen_offsets: Union[int, torch.Tensor] = 0,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
+    interleaved=False,
+    inplace=False,
+    conjugate=False,
+) -> torch.Tensor:
+    """
+    Arguments:
+        x: (batch, seqlen, nheads, headdim) if cu_seqlens is None
+            else (total_seqlen, nheads, headdim).
+        cos: (seqlen_ro, rotary_dim / 2)
+        sin: (seqlen_ro, rotary_dim / 2)
+        seqlen_offsets: integer or integer tensor of size (batch,)
+        cu_seqlens: (batch + 1,) or None
+        max_seqlen: int
+    Returns:
+        y: (batch, seqlen, nheads, headdim)
+    """
+    is_varlen = cu_seqlens is not None
+    if not is_varlen:
+        batch, seqlen, nheads, headdim = x.shape
+    else:
+        assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed"
+        total_seqlen, nheads, headdim = x.shape
+        batch_p_1 = cu_seqlens.shape[0]
+        batch = batch_p_1 - 1
+        seqlen = max_seqlen
+    seqlen_ro, rotary_dim = cos.shape
+    assert sin.shape == cos.shape
+    rotary_dim *= 2
+    assert rotary_dim <= headdim, f"rotary_dim must be <= headdim, but got {rotary_dim} and {headdim}"
+    assert headdim <= 256, "Only support headdim <= 256"
+    assert seqlen_ro >= seqlen, f"seqlen_ro must be >= seqlen, but got {seqlen_ro} and {seqlen}"
+
+    assert (
+        cos.dtype == sin.dtype
+    ), f"cos and sin must have the same dtype, got {cos.dtype} and {sin.dtype}"
+    assert (
+        x.dtype == cos.dtype
+    ), f"Input and cos/sin must have the same dtype, got {x.dtype} and {cos.dtype}"
+
+    cos, sin = cos.contiguous(), sin.contiguous()
+    if isinstance(seqlen_offsets, torch.Tensor):
+        assert seqlen_offsets.shape == (batch,)
+        assert seqlen_offsets.dtype in [torch.int32, torch.int64]
+        seqlen_offsets = seqlen_offsets.contiguous()
+    else:
+        assert seqlen_offsets + seqlen <= seqlen_ro
+
+    output = torch.empty_like(x) if not inplace else x
+    if rotary_dim < headdim and not inplace:
+        output[..., rotary_dim:].copy_(x[..., rotary_dim:])
+
+    BLOCK_K = (
+        32
+        if rotary_dim <= 32
+        else (64 if rotary_dim <= 64 else (128 if rotary_dim <= 128 else 256))
+    )
+    grid = lambda META: (triton.cdiv(seqlen, META["BLOCK_M"]), batch, nheads)  # noqa
+    BLOCK_M = 4 if interleaved else (8 if rotary_dim <= 128 else 4)
+
+    # Need this, otherwise Triton tries to launch from cuda:0 and we get
+    # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)
+    with torch.cuda.device(x.device.index):
+        rotary_kernel[grid](
+            output,  # data ptrs
+            x,
+            cos,
+            sin,
+            cu_seqlens,
+            seqlen_offsets,
+            seqlen,  # shapes
+            rotary_dim,
+            seqlen_ro,
+            output.stride(0) if not is_varlen else 0,  # batch_strides if not varlen else 0
+            output.stride(-3),  # seqlen_stride or total_seqlen_stride
+            output.stride(-2),  # nheads_stride
+            output.stride(-1),  # headdim_stride
+            x.stride(0) if not is_varlen else 0,  # batch_strides if not varlen else 0
+            x.stride(-3),  # seqlen stride or total_seqlen_stride
+            x.stride(-2),  # nheads stride
+            x.stride(-1),  # headdim stride
+            BLOCK_K,
+            isinstance(seqlen_offsets, torch.Tensor),
+            is_varlen,
+            interleaved,
+            conjugate,
+            BLOCK_M,
+        )
+    return output
+
+class ApplyRotaryEmb(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x,
+        cos,
+        sin,
+        interleaved=False,
+        inplace=False,
+        seqlen_offsets: Union[int, torch.Tensor] = 0,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+    ):
+        out = apply_rotary(
+            x,
+            cos,
+            sin,
+            seqlen_offsets=seqlen_offsets,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            interleaved=interleaved,
+            inplace=inplace,
+        )
+        if isinstance(seqlen_offsets, int):
+            ctx.save_for_backward(cos, sin, cu_seqlens)  # Can't save int with save_for_backward
+            ctx.seqlen_offsets = seqlen_offsets
+        else:
+            ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets)
+            ctx.seqlen_offsets = None
+        ctx.interleaved = interleaved
+        ctx.inplace = inplace
+        ctx.max_seqlen = max_seqlen
+        return out if not inplace else x
+
+    @staticmethod
+    def backward(ctx, do):
+        seqlen_offsets = ctx.seqlen_offsets
+        if seqlen_offsets is None:
+            cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors
+        else:
+            cos, sin, cu_seqlens = ctx.saved_tensors
+        # TD [2023-09-02]: For some reason Triton (2.0.0.post1) errors with
+        # "[CUDA]: invalid device context", and cloning makes it work. Idk why. Triton 2.1.0 works.
+        if not ctx.interleaved and not ctx.inplace:
+            do = do.clone()
+        dx = apply_rotary(
+            do,
+            cos,
+            sin,
+            seqlen_offsets=seqlen_offsets,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=ctx.max_seqlen,
+            interleaved=ctx.interleaved,
+            inplace=ctx.inplace,
+            conjugate=True,
+        )
+        return dx, None, None, None, None, None, None, None
+
+
+def apply_rotary_emb_triton(
+    x,
+    cos,
+    sin,
+    interleaved=False,
+    inplace=False,
+    seqlen_offsets: Union[int, torch.Tensor] = 0,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
+):
+    """
+    Arguments:
+        x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None
+            else (total_seqlen, nheads, headdim)
+        cos, sin: (seqlen_rotary, rotary_dim / 2)
+        interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
+            of 1st half and 2nd half (GPT-NeoX style).
+        inplace: if True, apply rotary embedding in-place.
+        seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount.
+            Most commonly used in inference when we have KV cache.
+        cu_seqlens: (batch + 1,) or None
+        max_seqlen: int
+    Return:
+        out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None
+            else (total_seqlen, nheads, headdim)
+    rotary_dim must be <= headdim
+    Apply rotary embedding to the first rotary_dim of x.
+    """
+    return ApplyRotaryEmb.apply(
+        x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen
+    )
+
+if torch.cuda.is_available() and torch.version.hip:
+    # do something specific for HIP
+    apply_rotary_emb = apply_rotary_emb_torch
+elif torch.cuda.is_available() and torch.version.cuda:
+    # do something specific for CUDA
+    apply_rotary_emb = apply_rotary_emb_triton
+else:
+    # do something for CPU
+    apply_rotary_emb = apply_rotary_emb_torch

ReSA/llm/kernel/tilelang_attention_with_kv_cache.py

@@ -0,0 +1,379 @@
+# Copyright (c) Tile-AI Corporation.
+# Licensed under the MIT License.
+
+import torch
+import torch.nn.functional as F
+import tilelang
+from tilelang.autotuner import *
+import tilelang.language as T
+import argparse
+import time
+import math
+
+def num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, num_m_blocks, size_one_kv_head,
+                         is_causal_or_local, max_splits):
+    """
+    Determines the optimal number of splits for maximizing GPU occupancy while balancing memory efficiency.
+
+    Parameters:
+    - total_mblocks (int): Total number of m_blocks.
+    - num_SMs (int): Number of Streaming Multiprocessors (SMs) in the GPU.
+    - num_n_blocks (int): Number of n_blocks.
+    - num_m_blocks (int): Number of m_blocks.
+    - size_one_kv_head (int): Size of one KV head in bytes.
+    - is_causal_or_local (bool): Indicates whether the operation is causal or local.
+    - max_splits (int): Maximum number of allowed splits.
+
+    Returns:
+    - int: The optimal number of splits.
+    """
+    # If we have enough m_blocks to almost fill the SMs, prefer 1 split unless memory constraints apply.
+    if total_mblocks >= 0.8 * num_SMs:
+        size_l2 = 50 * 1024 * 1024  # L2 cache size assumption (50MB)
+        # Only split if each KV head is too large for L2 and there are enough m_blocks
+        if size_one_kv_head > size_l2 and num_m_blocks >= num_SMs * 2 and not is_causal_or_local:
+            return min((size_one_kv_head + size_l2 - 1) // size_l2, max_splits)
+        else:
+            return 1
+
+    # If num_n_blocks is too small, we don't split
+    if num_n_blocks <= 4:
+        return 1
+
+    # Limit max_splits to a reasonable range
+    max_splits = min(max_splits, num_SMs, num_n_blocks)
+
+    max_efficiency = 0.0
+    efficiency = []
+
+    # Compute efficiency for different splits
+    for num_splits in range(1, max_splits + 1):
+        n_waves = (total_mblocks * num_splits) / num_SMs
+        eff = n_waves / math.ceil(n_waves)
+        # Track max efficiency
+        if eff > max_efficiency:
+            max_efficiency = eff
+
+        efficiency.append(eff)
+
+    # Find the smallest number of splits that achieves at least 85% of max efficiency
+    for num_splits in range(1, max_splits + 1):
+        if efficiency[num_splits - 1] >= 0.95 * max_efficiency:
+            return num_splits
+
+    return 1
+
+def flashattn(heads, heads_kv, dim, dim_v):
+    scale = (1.0 / dim)**0.5 * 1.44269504  # log2(e)
+    dtype = "bfloat16"
+    accum_dtype = "float"
+    kv_group_num = heads // heads_kv
+
+
+    def kernel_func(batch, block_N, block_H, block_M, num_split, num_stages, threads, seqlen_q, max_cache_seqlen):
+        shape_q = [batch, seqlen_q, heads, dim]
+        shape_k = [batch, max_cache_seqlen, heads_kv, dim]
+        shape_v = [batch, max_cache_seqlen, heads_kv, dim_v]
+        shape_o = [batch, seqlen_q, heads, dim_v]
+        part_shape = [batch, seqlen_q, heads, num_split, dim_v]
+        num_block_M = (seqlen_q + block_M - 1) // block_M
+
+        @T.macro
+        def flash_attn_split(
+                Q: T.Tensor(shape_q, dtype),
+                K: T.Tensor(shape_k, dtype),
+                V: T.Tensor(shape_v, dtype),
+                cache_seqlens: T.Tensor([batch], "int32"),
+                glse: T.Tensor([batch, seqlen_q, heads, num_split], accum_dtype),
+                Output_partial: T.Tensor(part_shape, accum_dtype),
+        ):
+            with T.Kernel(
+                    batch * num_block_M, heads_kv, num_split, threads=threads) as (bx, by, bz):
+                Q_shared = T.alloc_shared([block_M * block_H, dim], dtype)
+                K_shared = T.alloc_shared([block_N, dim], dtype)
+                V_shared = T.alloc_shared([block_N, dim_v], dtype)
+                acc_s = T.alloc_fragment([block_M * block_H, block_N], accum_dtype)
+                acc_s_cast = T.alloc_fragment([block_M * block_H, block_N], dtype)
+                acc_o = T.alloc_fragment([block_M * block_H, dim_v], accum_dtype)
+
+                scores_max = T.alloc_fragment([block_M * block_H], accum_dtype)
+                scores_max_prev = T.alloc_fragment([block_M * block_H], accum_dtype)
+                scores_scale = T.alloc_fragment([block_M * block_H], accum_dtype)
+                scores_sum = T.alloc_fragment([block_M * block_H], accum_dtype)
+                logsum = T.alloc_fragment([block_M * block_H], accum_dtype)
+
+                bid = T.floordiv(bx, num_block_M)
+                mid = T.floormod(bx, num_block_M) * block_M
+                hid = by
+                sid = bz
+
+                for i, d in T.Parallel(block_M * block_H, dim):
+                    i_m = T.floordiv(i, block_H)
+                    i_h = T.floormod(i, block_H)
+                    if i_h < kv_group_num:
+                        Q_shared[i, d] = Q[bid, mid + i_m, hid * kv_group_num + i_h, d]
+
+                # T.copy(Q[bid, mid:(mid + block_M), hid * kv_group_num : (hid + 1) * kv_group_num, :], Q_shared)
+                T.fill(acc_o, 0)
+                T.fill(logsum, 0)
+                T.fill(scores_max, -T.infinity(accum_dtype))
+                # num_blocks = actual_num_blocks[bid]
+                num_blocks = (cache_seqlens[bid] + block_N - 1) // block_N
+                blocks_per_split = T.floordiv(num_blocks, num_split)
+                remaining_blocks = T.floormod(num_blocks, num_split)
+                loop_range = (blocks_per_split + T.if_then_else(sid < remaining_blocks, 1, 0))
+                start = blocks_per_split * sid + T.min(sid, remaining_blocks)
+                for k in T.Pipelined(loop_range, num_stages=num_stages):
+                    i_s = start + k
+                    T.copy(
+                        K[bid, i_s * block_N: (i_s + 1) * block_N, hid, :], K_shared)
+                    T.clear(acc_s)
+                    T.gemm(
+                        Q_shared,
+                        K_shared,
+                        acc_s,
+                        transpose_B=True,
+                        policy=T.GemmWarpPolicy.FullRow)
+                    for i, j in T.Parallel(block_M * block_H, block_N):
+                        i_m = T.floordiv(i, block_H)
+                        i_h = T.floormod(i, block_H)
+                        acc_s[i_m * block_H + i_h, j] = T.if_then_else(i_s * block_N + j > cache_seqlens[bid] - seqlen_q + (mid + i_m), -T.infinity(accum_dtype), acc_s[i_m * block_H + i_h, j])
+                    T.copy(scores_max, scores_max_prev)
+                    T.fill(scores_max, -T.infinity(accum_dtype))
+                    T.reduce_max(acc_s, scores_max, dim=1, clear=False)
+                    for i in T.Parallel(block_M * block_H):
+                        scores_max[i] = T.if_then_else(scores_max[i] > scores_max_prev[i], scores_max[i], scores_max_prev[i])
+                        scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
+                    for i, j in T.Parallel(block_M * block_H, block_N):
+                        acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
+                    T.reduce_sum(acc_s, scores_sum, dim=1)
+                    for i in T.Parallel(block_M * block_H):
+                        logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
+                    T.copy(acc_s, acc_s_cast)
+                    for i, j in T.Parallel(block_M * block_H, dim_v):
+                        acc_o[i, j] *= scores_scale[i]
+                    T.copy(
+                        V[bid, i_s * block_N: (i_s + 1) * block_N, hid, :], V_shared)
+                    T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
+
+                for i, j in T.Parallel(block_M * block_H, dim_v):
+                    acc_o[i, j] /= logsum[i]
+                
+                for i in T.Parallel(block_M * block_H):
+                    logsum[i] = T.log2(logsum[i]) + scores_max[i] * scale
+                
+                for i in T.Parallel(block_M * block_H):
+                    i_m = T.floordiv(i, block_H)
+                    i_h = T.floormod(i, block_H)
+                    if i_h < kv_group_num:
+                        glse[bid, mid + i_m, hid * kv_group_num + i_h, sid] = logsum[i]
+
+                for i, v in T.Parallel(block_M * block_H, dim_v):
+                    i_m = T.floordiv(i, block_H)
+                    i_h = T.floormod(i, block_H)
+                    if i_h < kv_group_num:
+                        Output_partial[bid, mid + i_m, hid * kv_group_num + i_h, sid, v] = acc_o[i, v]
+                              
+        @T.macro
+        def combine(
+                glse: T.Tensor([batch, seqlen_q, heads, num_split], accum_dtype),
+                Output_partial: T.Tensor(part_shape, accum_dtype),
+                Output: T.Tensor(shape_o, dtype),
+        ):
+            with T.Kernel(heads, seqlen_q, batch, threads=128) as (bx, by, bz):
+                po_local = T.alloc_fragment([dim_v], accum_dtype)
+                o_accum_local = T.alloc_fragment([dim_v], accum_dtype)
+                lse_local_split = T.alloc_local([1], accum_dtype)
+                lse_logsum_local = T.alloc_local([1], accum_dtype)
+                lse_max_local = T.alloc_local([1], accum_dtype)
+                scale_local = T.alloc_local([1], accum_dtype)
+                max_split = T.alloc_local([1], "int32")
+
+                T.annotate_layout({
+                    lse_logsum_local: T.Fragment(lse_logsum_local.shape, forward_thread_fn=lambda i: i),
+                })
+
+                T.clear(lse_logsum_local)
+                T.clear(o_accum_local)
+                lse_max_local[0] = -T.infinity(accum_dtype)
+                for k in T.serial(num_split):
+                    lse_local_split[0] = glse[bz, by, bx, k]
+                    if (lse_local_split[0] != 0):
+                        max_split[0] = k
+                        lse_max_local[0] = T.max(lse_max_local[0], glse[bz, by, bx, k])
+
+                for k in T.Pipelined(num_split, num_stages=1):
+                    if k <= max_split[0]:
+                        lse_local_split[0] = glse[bz, by, bx, k]
+                        lse_logsum_local[0] += T.exp2(lse_local_split[0] - lse_max_local[0])
+                lse_logsum_local[0] = T.log2(lse_logsum_local[0]) + lse_max_local[0]
+                for k in T.serial(num_split):
+                    if k <= max_split[0]:
+                        for i in T.Parallel(dim_v):
+                            po_local[i] = Output_partial[bz, by, bx, k, i]
+                        lse_local_split[0] = glse[bz, by, bx, k]
+                        scale_local[0] = T.exp2(lse_local_split[0] - lse_logsum_local[0])
+                        for i in T.Parallel(dim_v):
+                            o_accum_local[i] += po_local[i] * scale_local[0]
+                for i in T.Parallel(dim_v):
+                    Output[bz, by, bx, i] = o_accum_local[i]
+
+        
+        @T.prim_func
+        def main(
+                Q: T.Tensor(shape_q, dtype),
+                K: T.Tensor(shape_k, dtype),
+                V: T.Tensor(shape_v, dtype),
+                cache_seqlens: T.Tensor([batch], "int32"),
+                glse: T.Tensor([batch, seqlen_q, heads, num_split], accum_dtype),
+                Output_partial: T.Tensor(part_shape, accum_dtype),
+                Output: T.Tensor(shape_o, dtype),
+        ):
+            flash_attn_split(Q, K, V, cache_seqlens, glse, Output_partial)
+            combine(glse, Output_partial, Output)
+
+        return main
+
+    return kernel_func
+
+
+class AttentionWithKVCache(torch.nn.Module):
+    def __init__(self, heads, heads_kv, dim, dim_v, seqlen_q):
+        super(AttentionWithKVCache, self).__init__()
+        self.heads = heads
+        self.heads_kv = heads_kv
+        self.dim = dim
+        self.dim_v = dim_v
+        self.block_N = 32
+        self.block_H = tilelang.next_power_of_2(heads // heads_kv)
+        self.block_M = seqlen_q
+
+        program = flashattn(heads, heads_kv, dim, dim_v)(
+            batch=T.symbolic("batch"),
+            block_N=self.block_N,
+            block_H=self.block_H,
+            block_M=self.block_M,
+            num_split=T.symbolic("num_split"),
+            num_stages=2,
+            threads=128,
+            seqlen_q=seqlen_q,
+            max_cache_seqlen=T.symbolic("max_cache_seqlen"),
+        )
+
+        self.kernel = tilelang.compile(
+            program,
+            out_idx=-1,
+            target='cuda',
+            execution_backend="cython"
+        )
+
+        props = torch.cuda.get_device_properties(torch.device("cuda:0"))
+        self.num_sm = props.multi_processor_count
+
+    def forward(self, query, key, value, cache_seqlens):
+        batch = query.shape[0]
+        seqlen_q = query.shape[1]
+        heads = self.heads
+        heads_kv = self.heads_kv
+        dim = self.dim
+        dim_v = self.dim_v
+
+        # Compute static scheduling parameters
+        num_m_blocks = (seqlen_q + self.block_M - 1) // self.block_M
+        num_n_blocks = (cache_seqlens.max().item() + self.block_N - 1) // self.block_N
+        size_one_kv_head = num_n_blocks * self.block_N * (dim + dim_v) * 2
+        total_mblocks = batch * heads_kv * num_m_blocks
+        # num_sm = 132
+        num_sm = self.num_sm
+
+        num_split = num_splits_heuristic(
+            total_mblocks, num_sm, num_n_blocks, num_m_blocks,
+            size_one_kv_head, is_causal_or_local=True, max_splits=16
+        )
+
+        glse = torch.empty((batch, seqlen_q, heads, num_split), dtype=torch.float32, device='cuda')
+        output_partial = torch.empty((batch, seqlen_q, heads, num_split, dim_v), dtype=torch.float32, device='cuda')
+
+        output = self.kernel(
+            query, key, value, cache_seqlens,
+            glse, output_partial
+        )
+        return output
+
+def ref_program_fa(query, key, value, cache_seqlens):
+    # latency reference
+    # from flash_attn_interface import flash_attn_with_kvcache, flash_attn_func # fa3
+    from flash_attn import flash_attn_with_kvcache, flash_attn_func #fa2
+    output = flash_attn_with_kvcache(query, key, value, cache_seqlens=cache_seqlens)
+    return output
+
+
+
+def debug(name,expect, actual, atol=1e-3, rtol=1e-3):
+    all_close = torch.allclose(expect, actual, atol=atol, rtol=rtol)
+    print(name + "  all_close={}".format(all_close))
+    if not all_close:
+        # print(expect[3, 28])
+        # print(actual[3, 28])
+        diff = (expect - actual).abs()
+        print("all_close={}, max={}, min={}, mean={}".format(all_close, diff.max().item(), diff.min().item(), diff.mean().item()))
+        max_indices = torch.nonzero(diff == diff.max().item())
+        first_index = tuple(max_indices[0].tolist())
+        print(f"Index: {first_index}, expect: {expect[first_index]}, actual: {actual[first_index]}") 
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=4, help='batch size')
+    parser.add_argument('--seqlen_q', type=int, default=32, help='sequence length')
+    parser.add_argument('--heads', type=int, default=28, help='heads')
+    parser.add_argument('--heads_kv', type=int, default=4, help='heads_kv')
+    parser.add_argument('--max_cache_seqlen', type=int, default=65536, help='kvcache sequence length')
+    parser.add_argument('--dim', type=int, default=128, help='dim')
+    parser.add_argument('--dim_v', type=int, default=128, help='dim_v')
+    parser.add_argument('--block_size', type=int, default=32, help='block_size')
+    parser.add_argument('--load_from_file', type=str, default=None, help='load from file')
+    args = parser.parse_args()
+
+    batch, seqlen_q, heads, heads_kv, max_cache_seqlen, dim, dim_v = args.batch, args.seqlen_q, args.heads, args.heads_kv, args.max_cache_seqlen, args.dim, args.dim_v
+    block_size = args.block_size
+
+    dtype = torch.bfloat16
+
+    Q = torch.randn((batch, seqlen_q, heads, dim), dtype=dtype, device='cuda')
+    K = torch.randn((batch, max_cache_seqlen, heads_kv, dim), dtype=dtype, device='cuda')
+    V = torch.randn((batch, max_cache_seqlen, heads_kv, dim_v), dtype=dtype, device='cuda')
+
+    cache_seqlens = torch.randint(max_cache_seqlen - 32, max_cache_seqlen, (batch,), device='cuda', dtype=torch.int32)
+    print("cache_seqlens: ", cache_seqlens)
+    # parity reference
+    ref = ref_program_fa(Q, K, V, cache_seqlens)
+    # ref = ref_program_triton(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, max_num_blocks, block_size)
+    # out = kernel(Q, K, V, block_indices, cache_seqlens, actual_num_blocks, glse, Output_partial)
+    # out = sparse_gqa_decode_varlen_indice(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, block_size)
+    sparse_kernel = AttentionWithKVCache(heads, heads_kv, dim, dim_v, seqlen_q)
+    out = sparse_kernel(Q, K, V, cache_seqlens)
+    debug("output", ref, out, atol=1e-3, rtol=1e-3)
+
+    ## latency reference
+    for i in range(10):
+        ref = ref_program_fa(Q, K, V, cache_seqlens)
+    torch.cuda.synchronize()
+    start = time.time()
+    for i in range(100):
+        ref = ref_program_fa(Q, K, V, cache_seqlens)
+    torch.cuda.synchronize()
+    print("dense time: ", (time.time() - start) / 100*1000)
+
+    for i in range(10):
+        # out = sparse_gqa_decode_varlen_indice(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, block_size)
+        out = sparse_kernel(Q, K, V, cache_seqlens)
+    torch.cuda.synchronize()
+    start = time.time()
+    for i in range(100):
+        # out = sparse_gqa_decode_varlen_indice(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, block_size)
+        out = sparse_kernel(Q, K, V, cache_seqlens)
+    torch.cuda.synchronize()
+    print("sparse time: ", (time.time() - start) / 100*1000)
+
+
+

ReSA/llm/kernel/tilelang_sparse_decoding.py

@@ -0,0 +1,538 @@
+# Copyright (c) Tile-AI Corporation.
+# Licensed under the MIT License.
+
+import torch
+import torch.nn.functional as F
+import tilelang
+from tilelang.autotuner import *
+import tilelang.language as T
+from einops import rearrange, einsum
+import argparse
+import time
+import math
+
+def num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, num_m_blocks, size_one_kv_head,
+                         is_causal_or_local, max_splits):
+    """
+    Determines the optimal number of splits for maximizing GPU occupancy while balancing memory efficiency.
+
+    Parameters:
+    - total_mblocks (int): Total number of m_blocks.
+    - num_SMs (int): Number of Streaming Multiprocessors (SMs) in the GPU.
+    - num_n_blocks (int): Number of n_blocks.
+    - num_m_blocks (int): Number of m_blocks.
+    - size_one_kv_head (int): Size of one KV head in bytes.
+    - is_causal_or_local (bool): Indicates whether the operation is causal or local.
+    - max_splits (int): Maximum number of allowed splits.
+
+    Returns:
+    - int: The optimal number of splits.
+    """
+    # If we have enough m_blocks to almost fill the SMs, prefer 1 split unless memory constraints apply.
+    if total_mblocks >= 0.8 * num_SMs:
+        size_l2 = 50 * 1024 * 1024  # L2 cache size assumption (50MB)
+        # Only split if each KV head is too large for L2 and there are enough m_blocks
+        if size_one_kv_head > size_l2 and num_m_blocks >= num_SMs * 2 and not is_causal_or_local:
+            return min((size_one_kv_head + size_l2 - 1) // size_l2, max_splits)
+        else:
+            return 1
+
+    # If num_n_blocks is too small, we don't split
+    if num_n_blocks <= 4:
+        return 1
+
+    # Limit max_splits to a reasonable range
+    max_splits = min(max_splits, num_SMs, num_n_blocks)
+
+    max_efficiency = 0.0
+    efficiency = []
+
+    # Compute efficiency for different splits
+    for num_splits in range(1, max_splits + 1):
+        n_waves = (total_mblocks * num_splits) / num_SMs
+        eff = n_waves / math.ceil(n_waves)
+        # Track max efficiency
+        if eff > max_efficiency:
+            max_efficiency = eff
+
+        efficiency.append(eff)
+
+    # Find the smallest number of splits that achieves at least 85% of max efficiency
+    for num_splits in range(1, max_splits + 1):
+        if efficiency[num_splits - 1] >= 0.95 * max_efficiency:
+            return num_splits
+
+    return 1
+
+def flashattn(heads, heads_kv, dim, dim_v):
+    scale = (1.0 / dim)**0.5 * 1.44269504  # log2(e)
+    dtype = "float16"
+    accum_dtype = "float"
+    kv_group_num = heads // heads_kv
+
+
+    def kernel_func(batch, block_N, block_H, num_split, num_stages, threads, max_cache_seqlen, max_selected_blocks):
+        shape_q = [batch, heads, dim]
+        shape_k = [batch, max_cache_seqlen, heads_kv, dim]
+        shape_v = [batch, max_cache_seqlen, heads_kv, dim_v]
+        shape_indices = [batch, heads_kv, max_selected_blocks]
+        shape_o = [batch, heads, dim_v]
+        part_shape = [batch, heads, num_split, dim_v]
+        valid_block_H = min(block_H, kv_group_num)
+
+        @T.macro
+        def flash_attn_split(
+                Q: T.Tensor(shape_q, dtype),
+                K: T.Tensor(shape_k, dtype),
+                V: T.Tensor(shape_v, dtype),
+                block_indices: T.Tensor(shape_indices, "int32"),
+                cache_seqlens: T.Tensor([batch], "int32"),
+                # actual_num_blocks: T.Tensor([batch], "int32"),
+                glse: T.Tensor([batch, heads, num_split], accum_dtype),
+                Output_partial: T.Tensor(part_shape, accum_dtype),
+        ):
+            with T.Kernel(
+                    batch, heads // valid_block_H, num_split, threads=threads) as (bx, by, bz):
+                Q_shared = T.alloc_shared([block_H, dim], dtype)
+                K_shared = T.alloc_shared([block_N, dim], dtype)
+                V_shared = T.alloc_shared([block_N, dim_v], dtype)
+                # O_shared = T.alloc_shared([valid_block_H, dim_v], dtype)
+                acc_s = T.alloc_fragment([block_H, block_N], accum_dtype)
+                acc_s_cast = T.alloc_fragment([block_H, block_N], dtype)
+                acc_o = T.alloc_fragment([block_H, dim_v], accum_dtype)
+
+                scores_max = T.alloc_fragment([block_H], accum_dtype)
+                scores_max_prev = T.alloc_fragment([block_H], accum_dtype)
+                scores_scale = T.alloc_fragment([block_H], accum_dtype)
+                scores_sum = T.alloc_fragment([block_H], accum_dtype)
+                logsum = T.alloc_fragment([block_H], accum_dtype)
+                has_valid_block=T.alloc_var("bool")
+                # num_blocks = T.alloc_local([1], "int32")
+
+                bid = bx
+                hid = by
+                sid = bz
+                cur_kv_head = hid // (kv_group_num // valid_block_H)
+
+                T.copy(Q[bid, hid * valid_block_H:hid * valid_block_H + block_H, :], Q_shared)
+                T.fill(acc_o, 0)
+                T.fill(logsum, 0)
+                T.fill(scores_max, -T.infinity(accum_dtype))
+                # num_blocks = actual_num_blocks[bid]
+                num_blocks = max_selected_blocks
+                blocks_per_split = T.floordiv(num_blocks, num_split)
+                remaining_blocks = T.floormod(num_blocks, num_split)
+                loop_range = (blocks_per_split + T.if_then_else(sid < remaining_blocks, 1, 0))
+                start = blocks_per_split * sid + T.min(sid, remaining_blocks)
+                has_valid_block=False
+                # if (start < num_blocks):
+                for k in T.Pipelined(loop_range, num_stages=num_stages):
+                    i_s = block_indices[bid, cur_kv_head, start + k] 
+                    if i_s >= 0:
+                        has_valid_block = True
+                        T.copy(
+                            K[bid, i_s * block_N: (i_s + 1) * block_N,
+                            cur_kv_head, :], K_shared)
+                        T.clear(acc_s)
+                        T.gemm(
+                            Q_shared,
+                            K_shared,
+                            acc_s,
+                            transpose_B=True,
+                            policy=T.GemmWarpPolicy.FullRow)
+                        # if k == 0: # assume block_indices is sorted in reverse order, otherwise, remove this if condition
+                        for i, j in T.Parallel(block_H, block_N):
+                            acc_s[i, j] = T.if_then_else(i_s * block_N + j >= cache_seqlens[bid], -T.infinity(accum_dtype), acc_s[i, j])
+                        T.copy(scores_max, scores_max_prev)
+                        T.fill(scores_max, -T.infinity(accum_dtype))
+                        T.reduce_max(acc_s, scores_max, dim=1, clear=False)
+                        for i in T.Parallel(block_H):
+                            scores_max[i] = T.if_then_else(scores_max[i] > scores_max_prev[i], scores_max[i], scores_max_prev[i])
+                            scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
+                        for i, j in T.Parallel(block_H, block_N):
+                            acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
+                        T.reduce_sum(acc_s, scores_sum, dim=1)
+                        for i in T.Parallel(block_H):
+                            logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
+                        T.copy(acc_s, acc_s_cast)
+                        for i, j in T.Parallel(block_H, dim_v):
+                            acc_o[i, j] *= scores_scale[i]
+                        T.copy(
+                            V[bid, i_s * block_N: (i_s + 1) * block_N,
+                            cur_kv_head, :], V_shared)
+                        T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
+                if has_valid_block:
+                    for i, j in T.Parallel(block_H, dim_v):
+                        acc_o[i, j] /= logsum[i]
+                    
+                    for i in T.Parallel(block_H):
+                        logsum[i] = T.log2(logsum[i]) + scores_max[i] * scale
+                
+                for i in T.Parallel(block_H):
+                    if i < valid_block_H:
+                        glse[bid, hid * valid_block_H + i, sid] = logsum[i]
+
+                for i, j in T.Parallel(block_H, dim_v):
+                    if i < valid_block_H:
+                        Output_partial[bid, hid * valid_block_H + i, sid, j] = acc_o[i, j]
+                              
+        @T.macro
+        def combine(
+                glse: T.Tensor([batch, heads, num_split], accum_dtype),
+                Output_partial: T.Tensor(part_shape, accum_dtype),
+                Output: T.Tensor(shape_o, dtype),
+        ):
+            with T.Kernel(heads, batch, threads=128) as (by, bz):
+                po_local = T.alloc_fragment([dim_v], accum_dtype)
+                o_accum_local = T.alloc_fragment([dim_v], accum_dtype)
+                lse_local_split = T.alloc_local([1], accum_dtype)
+                lse_logsum_local = T.alloc_local([1], accum_dtype)
+                lse_max_local = T.alloc_local([1], accum_dtype)
+                scale_local = T.alloc_local([1], accum_dtype)
+                max_split = T.alloc_local([1], "int32")
+
+                T.annotate_layout({
+                    lse_logsum_local: T.Fragment(lse_logsum_local.shape, forward_thread_fn=lambda i: i),
+                })
+
+                T.clear(lse_logsum_local)
+                T.clear(o_accum_local)
+                lse_max_local[0] = -T.infinity(accum_dtype)
+                for k in T.serial(num_split):
+                    lse_local_split[0] = glse[bz, by, k]
+                    if (lse_local_split[0] != 0):
+                        max_split[0] = k
+                        lse_max_local[0] = T.max(lse_max_local[0], glse[bz, by, k])
+
+                for k in T.Pipelined(num_split, num_stages=1):
+                    if k <= max_split[0]:
+                        lse_local_split[0] = glse[bz, by, k]
+                        lse_logsum_local[0] += T.exp2(lse_local_split[0] - lse_max_local[0])
+                lse_logsum_local[0] = T.log2(lse_logsum_local[0]) + lse_max_local[0]
+                for k in T.serial(num_split):
+                    if k <= max_split[0]:
+                        for i in T.Parallel(dim_v):
+                            po_local[i] = Output_partial[bz, by, k, i]
+                        lse_local_split[0] = glse[bz, by, k]
+                        scale_local[0] = T.exp2(lse_local_split[0] - lse_logsum_local[0])
+                        for i in T.Parallel(dim_v):
+                            o_accum_local[i] += po_local[i] * scale_local[0]
+                for i in T.Parallel(dim_v):
+                    Output[bz, by, i] = o_accum_local[i]
+
+        
+        @T.prim_func
+        def main(
+                Q: T.Tensor(shape_q, dtype),
+                K: T.Tensor(shape_k, dtype),
+                V: T.Tensor(shape_v, dtype),
+                block_indices: T.Tensor(shape_indices, "int32"),
+                cache_seqlens: T.Tensor([batch], "int32"),
+                # actual_num_blocks: T.Tensor([batch], "int32"),
+                glse: T.Tensor([batch, heads, num_split], accum_dtype),
+                Output_partial: T.Tensor(part_shape, accum_dtype),
+                Output: T.Tensor(shape_o, dtype),
+        ):
+            # flash_attn_split(Q, K, V, block_indices, cache_seqlens, actual_num_blocks, glse, Output_partial)
+            flash_attn_split(Q, K, V, block_indices, cache_seqlens, glse, Output_partial)
+            combine(glse, Output_partial, Output)
+
+        return main
+
+    return kernel_func
+
+
+class SparseFlashAttn(torch.nn.Module):
+    def __init__(self, heads, heads_kv, dim, dim_v, block_size):
+        super(SparseFlashAttn, self).__init__()
+        self.heads = heads
+        self.heads_kv = heads_kv
+        self.dim = dim
+        self.dim_v = dim_v
+        self.block_size = block_size
+
+        self.block_H = 64
+
+        program = flashattn(heads, heads_kv, dim, dim_v)(
+            batch=T.symbolic("batch"),
+            block_N=block_size,
+            block_H=self.block_H,
+            num_split=T.symbolic("num_split"),
+            num_stages=2,
+            threads=128,
+            max_cache_seqlen=T.symbolic("max_cache_seqlen"),
+            max_selected_blocks=T.symbolic("max_selected_blocks")
+        )
+
+        self.kernel = tilelang.compile(
+            program,
+            out_idx=-1,
+            target='cuda',
+            execution_backend="cython"
+        )
+
+        props = torch.cuda.get_device_properties(torch.device("cuda:0"))
+        self.num_sm = props.multi_processor_count
+
+    def forward(self, query, key, value, block_indices, cache_seqlens):
+        batch = query.shape[0]
+        heads = self.heads
+        heads_kv = self.heads_kv
+        dim = self.dim
+        dim_v = self.dim_v
+        block_size = self.block_size
+        max_selected_blocks = block_indices.shape[-1]
+
+        # Compute static scheduling parameters
+        num_m_blocks = 1 * (heads // heads_kv + self.block_H - 1) // self.block_H
+        num_n_blocks = max_selected_blocks
+        size_one_kv_head = max_selected_blocks * block_size * (dim + dim_v) * 2
+        total_mblocks = batch * heads_kv * num_m_blocks
+        # num_sm = 132
+        num_sm = self.num_sm
+
+        num_split = num_splits_heuristic(
+            total_mblocks, num_sm, num_n_blocks, num_m_blocks,
+            size_one_kv_head, is_causal_or_local=True, max_splits=16
+        )
+
+        # Function to compile
+        # def compute_actual_num_blocks(block_indices):
+        #     actual_num_blocks = torch.sum(block_indices != -1, dim=-1).to(torch.int32)
+        #     actual_num_blocks = actual_num_blocks[:, 0]  # [batch]
+        #     return actual_num_blocks
+        # compiled_fn = torch.compile(compute_actual_num_blocks)
+        # actual_num_blocks = compiled_fn(block_indices)
+        glse = torch.empty((batch, heads, num_split), dtype=torch.float32, device='cuda')
+        output_partial = torch.empty((batch, heads, num_split, dim_v), dtype=torch.float32, device='cuda')
+
+        
+        # output = self.kernel(
+        #     query, key, value, block_indices, cache_seqlens,
+        #     actual_num_blocks, glse, output_partial
+        # )
+        output = self.kernel(
+            query, key, value, block_indices, cache_seqlens,
+            glse, output_partial
+        )
+        return output
+
+def sparse_gqa_decode_varlen_indice(query, key, value, block_indices, cache_seqlens, max_cache_seqlen, block_size):
+    """
+    Args:
+        query: [batch, heads, dim]
+        key: [batch, max_cache_seqlen, heads_kv, dim]
+        value: [batch, max_cache_seqlen, heads_kv, dim_v]
+        block_indices: [batch, heads_kv, max_selected_blocks], indices of selected blocks, -1 for padding
+        cache_seqlens: [batch], sequence lengths of the kvcache
+        max_cache_seqlen: maximum sequence length of kvcache
+        block_size: block size
+    Returns:
+        output: [batch, heads, dim_v]
+
+    """
+    
+    batch, heads, dim = query.shape
+    heads_kv = key.shape[2]
+    dim_v = value.shape[-1]
+    max_selected_blocks = block_indices.shape[-1]
+    block_H = 64
+
+
+    actual_num_blocks = torch.sum(block_indices != -1, dim=-1).to(torch.int32)
+    actual_num_blocks = actual_num_blocks[:,0] #[batch],  number of valid blocks, assum all groups in the same batch have the same number of blocks
+   
+    # get num_split
+    num_m_blocks = 1 * (heads // heads_kv + block_H - 1) // block_H
+    num_n_blocks = max_selected_blocks#(kv_seqlen  + block_size - 1 ) // block_size
+    # num_n_blocks = torch.sum(actual_num_blocks, dim=-1).item() * heads_kv # total number of blocks
+
+    size_one_kv_head = max_selected_blocks * block_size * (dim + dim_v) * 2 #kv_seqlen * (dim + dim_v) * 2
+    total_mblocks = batch * heads_kv * num_m_blocks
+    num_sm = 132
+    num_split = num_splits_heuristic(total_mblocks, num_sm, num_n_blocks, num_m_blocks, size_one_kv_head, is_causal_or_local=True, max_splits=128)
+   
+    program = flashattn(
+        batch, heads, heads_kv, dim, dim_v)(
+                block_N=block_size, block_H=block_H, num_split=T.symbolic("num_split"), num_stages=2, threads=128, 
+                max_cache_seqlen=T.symbolic("max_cache_seqlen"), max_selected_blocks=T.symbolic("max_selected_blocks"))
+
+    glse = torch.empty((batch, heads, num_split), dtype=torch.float32, device='cuda')
+    Output_partial = torch.empty((batch, heads, num_split, dim_v), dtype=torch.float32, device='cuda')
+    kernel = tilelang.compile(program, out_idx=-1, target='cuda', execution_backend="cython", pass_configs={"tl.config_index_bitwidth": 64})
+    # print(kernel.get_kernel_source())
+
+    # output = kernel(query, key, value, block_indices, cache_seqlens, actual_num_blocks, glse, Output_partial)
+    output = kernel(query, key, value, block_indices, cache_seqlens, glse, Output_partial)
+    return output
+
+def ref_program_torch(query, key, value,  block_indices, cache_seqlens, max_cache_seqlen, num_blocks, block_size):
+
+    batch, heads, dim = query.shape
+    heads_kv = key.shape[2]
+    dim_v = value.shape[-1]
+    num_head_groups = query.shape[1] // key.shape[2]
+    scale = dim**0.5
+    key = rearrange(key, 'b n h d -> b h n d')  # [batch_size, heads_kv, seqlen_kv, dim]
+    value = rearrange(value, 'b n h d -> b h n d')  # [batch_size, heads_kv, seqlen_kv, dim]
+
+    query = rearrange(
+        query, 'b (h g) d -> b g h d',
+        g=num_head_groups)  # [batch_size, num_head_groups, heads_kv, dim]
+
+    scores = einsum(
+        query, key,
+        'b g h d, b h s d -> b g h s')  # [batch_size, num_head_groups, heads_kv, seqlen_kv]
+
+    sparse_mask = torch.zeros_like(scores)
+    # Assign mask values based on block_indices
+    for b in range(batch):
+        for h in range(heads_kv):
+            valid_indices = block_indices[b, h]  # Extract indices for this batch and head
+            for idx in valid_indices:
+                if idx >= 0:
+                    sparse_mask[b, :, h, idx * block_size: (idx + 1) * block_size] = 1
+    scores = scores.masked_fill(sparse_mask == 0, float('-inf'))
+    
+
+    range_len = torch.arange(scores.shape[-1], device='cuda').unsqueeze(0)
+    cache_seqlens_expanded = cache_seqlens.unsqueeze(1)
+    pad_mask = range_len >= cache_seqlens_expanded     
+    pad_mask = pad_mask[:, None, None, :]
+    scores = scores.masked_fill(pad_mask, float('-inf'))
+    attention = F.softmax(
+        scores / scale, dim=-1)  # [batch_size, num_head_groups, heads_kv, seqlen_kv]
+
+    out = einsum(attention, value,
+                 'b g h s, b h s d -> b g h d')  # [batch_size, num_head_groups, heads_kv, dim]
+    out = rearrange(out, 'b g h d -> b (h g) d')  # [batch_size, heads, dim]
+    return out
+
+
+def ref_program_fa(query, key, value,  block_indices, cache_seqlens, max_cache_seqlen, num_blocks, block_size):
+    # latency reference
+    # from flash_attn_interface import flash_attn_with_kvcache, flash_attn_func # fa3
+    from flash_attn import flash_attn_with_kvcache, flash_attn_func #fa2
+    query = query.unsqueeze(1)
+    output = flash_attn_with_kvcache(query, key, value, cache_seqlens=cache_seqlens)
+    output = output.squeeze(1)
+    return output
+
+
+
+def debug(name,expect, actual, atol=1e-3, rtol=1e-3):
+    all_close = torch.allclose(expect, actual, atol=atol, rtol=rtol)
+    print(name + "  all_close={}".format(all_close))
+    if not all_close:
+        # print(expect[3, 28])
+        # print(actual[3, 28])
+        diff = (expect - actual).abs()
+        print("all_close={}, max={}, min={}, mean={}".format(all_close, diff.max().item(), diff.min().item(), diff.mean().item()))
+        max_indices = torch.nonzero(diff == diff.max().item())
+        first_index = tuple(max_indices[0].tolist())
+        print(f"Index: {first_index}, expect: {expect[first_index]}, actual: {actual[first_index]}") 
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch', type=int, default=4, help='batch size')
+    parser.add_argument('--heads', type=int, default=28, help='heads')
+    parser.add_argument('--heads_kv', type=int, default=4, help='heads_kv')
+    parser.add_argument('--max_cache_seqlen', type=int, default=1048576, help='kvcache sequence length')
+    parser.add_argument('--dim', type=int, default=128, help='dim')
+    parser.add_argument('--dim_v', type=int, default=128, help='dim_v')
+    parser.add_argument('--sparse_ratio', type=float, default=0.9, help='sparse ratio')
+    parser.add_argument('--block_size', type=int, default=32, help='block_size')
+    parser.add_argument('--load_from_file', type=str, default=None, help='load from file')
+    args = parser.parse_args()
+
+    block_H = 64
+    if args.load_from_file is None:
+        batch, heads, heads_kv, max_cache_seqlen, dim, dim_v = args.batch, args.heads, args.heads_kv, args.max_cache_seqlen, args.dim, args.dim_v
+        sparse_ratio = args.sparse_ratio
+        block_size = args.block_size
+
+        max_selected_blocks = int(math.ceil(max_cache_seqlen * (1-sparse_ratio)/ block_size))
+        print("max_selected_blocks: ", max_selected_blocks)
+        dtype = torch.float16
+
+        Q = torch.randn((batch, heads, dim), dtype=dtype, device='cuda')
+        K = torch.randn((batch, max_cache_seqlen, heads_kv, dim), dtype=dtype, device='cuda')
+        V = torch.randn((batch, max_cache_seqlen, heads_kv, dim_v), dtype=dtype, device='cuda')
+        cache_seqlens = torch.full((batch,), max_cache_seqlen, dtype=torch.int32, device='cuda')
+        # cache_seqlens = torch.randint(1, max_cache_seqlen, (batch,), dtype=torch.int32, device='cuda')
+        # cache_seqlens = torch.full((batch,), max_cache_seqlen, dtype=torch.int32, device='cuda')
+        # Ensure at least one element equals cache_seqlen
+        random_index = torch.randint(0, batch, (1,), device='cuda').item()  # Select a random index
+        cache_seqlens[random_index] = max_cache_seqlen  # Assign cache_seqlen to ensure at least one occurrence
+    else:
+        save_dict = torch.load(args.load_from_file)
+        Q = save_dict["xq"]
+        K = save_dict["key"]
+        V = save_dict["value"]
+        block_indices = save_dict["sparse_indices"]
+        cache_seqlens = save_dict["seqlens_k"]
+        batch, heads, dim = Q.shape
+        heads_kv = K.shape[2]
+        max_cache_seqlen = K.shape[1]
+        dim_v = V.shape[-1]
+        block_size = max_cache_seqlen // block_indices.shape[-1]
+        max_selected_blocks = block_indices.shape[-1]
+        print(f"Load debug data from file with batch={batch}, heads={heads}, heads_kv={heads_kv}, max_cache_seqlen={max_cache_seqlen}, dim={dim}, dim_v={dim_v}, block_size={block_size}")
+
+    print("cache_seqlens: ", cache_seqlens)
+
+    max_valid_num_blocks = torch.ceil(cache_seqlens / block_size).int()
+    print("max_valid_num_blocks: ", max_valid_num_blocks)
+    # Initialize block_indices with -1 (for padding blocks)
+    block_indices = torch.full((batch, heads_kv, max_selected_blocks), -1, dtype=torch.int32, device='cuda')
+
+    # Assign valid indices while ensuring no duplicates within each batch-group
+    for b in range(batch):
+        max_valid_block = max_valid_num_blocks[b].item()  # Max valid blocks for this batch
+        if max_valid_block > 0:  # Ensure there's at least one valid block
+            for h in range(heads_kv):
+                valid_indices = torch.randperm(max_valid_block, device='cuda', dtype=torch.int32)[:max_selected_blocks]
+                block_indices[b, h, :len(valid_indices)] = valid_indices 
+
+    # Sort indices within each batch-group for consistency
+    block_indices, _ = block_indices.sort(dim=-1, descending=True)
+    # print("block_indices: ", block_indices)
+    actual_num_blocks = torch.sum(block_indices != -1, dim=-1).to(torch.int32)[:,0]
+    print("actual_num_blocks: ", actual_num_blocks)
+    # print(block_indices.shape, actual_num_blocks.shape)
+   
+    max_num_blocks = torch.max(max_valid_num_blocks).item()
+    print("max_num_blocks: ", max_num_blocks)
+
+    # parity reference
+    ref = ref_program_torch(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, max_num_blocks, block_size)
+    # ref = ref_program_triton(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, max_num_blocks, block_size)
+    # out = kernel(Q, K, V, block_indices, cache_seqlens, actual_num_blocks, glse, Output_partial)
+    # out = sparse_gqa_decode_varlen_indice(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, block_size)
+    sparse_kernel = SparseFlashAttn(heads, heads_kv, dim, dim_v, block_size)
+    out = sparse_kernel(Q, K, V, block_indices, cache_seqlens)
+    debug("output", ref, out, atol=1e-3, rtol=1e-3)
+
+    ## latency reference
+    for i in range(10):
+        ref = ref_program_fa(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, max_num_blocks, block_size)
+    torch.cuda.synchronize()
+    start = time.time()
+    for i in range(100):
+        ref = ref_program_fa(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, max_num_blocks, block_size)
+    torch.cuda.synchronize()
+    print("dense time: ", (time.time() - start) / 100*1000)
+
+    for i in range(10):
+        # out = sparse_gqa_decode_varlen_indice(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, block_size)
+        out = sparse_kernel(Q, K, V, block_indices, cache_seqlens)
+    torch.cuda.synchronize()
+    start = time.time()
+    for i in range(100):
+        # out = sparse_gqa_decode_varlen_indice(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen, block_size)
+        out = sparse_kernel(Q, K, V, block_indices, cache_seqlens)
+    torch.cuda.synchronize()
+    print("sparse time: ", (time.time() - start) / 100*1000)
+
+
+