NVFP4 QSA sparse attention (#408) spills under Triton 3.8: ~10x slower prefill and ~550 MiB local memory #675

"""QSA sparse attention (qwen4_exp prefill shapes, nvfp4 KV): registers, spills, local-memory growth and time
for one launch-config override.  usage: qsa_cfg.py WARPS|- MAXNREG|-   (fresh process per config: lmem is sticky)"""
import sys
import torch
import triton
from freetoken.kernel.triton.qsa import attend

warps = None if sys.argv[1] == "-" else int(sys.argv[1])
maxnreg = None if sys.argv[2] == "-" else int(sys.argv[2])
mode = sys.argv[3] if len(sys.argv) > 3 else "nvfp4"
orig = attend._qsa_sparse_paged_gqa_splitk_kernel


class Proxy:
    def __getitem__(self, grid):
        launch = orig[grid]

        def run(*a, **kw):
            if warps is not None and kw.get("BLOCK_N") == 64:
                kw["num_warps"] = warps
            if maxnreg is not None and kw.get("BLOCK_N") == 64:
                kw["maxnreg"] = maxnreg
            return launch(*a, **kw)
        return run


attend._qsa_sparse_paged_gqa_splitk_kernel = Proxy()
dev = "cuda"
torch.zeros(1, device=dev)


def outside():
    torch.cuda.synchronize()
    f, t = torch.cuda.mem_get_info()
    return (t - f - torch.cuda.memory_reserved()) / 2**20


D, KVH, QH, PAGE, TOPK = 256, 2, 24, 64, 2112
pages = 80
slots = pages * PAGE
k = torch.randn(slots, KVH * D, device=dev, dtype=torch.bfloat16)
v = torch.randn_like(k) * 0.25
if mode == "nvfp4":
    from freetoken.kernel.triton.kv_nvfp4 import quantize_nvfp4_to_cache
    codes = [torch.zeros((slots, KVH, D // 2), dtype=torch.uint8, device=dev) for _ in range(2)]
    scales = [torch.zeros((slots, KVH), dtype=torch.float32, device=dev) for _ in range(2)]
    blocks = [torch.zeros((slots, KVH, D // 16), dtype=torch.uint8, device=dev) for _ in range(2)]
    quantize_nvfp4_to_cache(k, v, torch.arange(slots, dtype=torch.int32, device=dev), codes[0], codes[1], scales[0], scales[1], blocks[0], blocks[1])
    packed = (pages, PAGE, KVH, D // 2)
    kv = (codes[0].view(packed), codes[1].view(packed))
    extra = dict(k_scale=scales[0], v_scale=scales[1], kv_quant="nvfp4", k_block_scale=blocks[0], v_block_scale=blocks[1])
else:
    kv = (k.view(pages, PAGE, KVH, D), v.view(pages, PAGE, KVH, D))
    extra = {}
block_table = torch.arange(pages, dtype=torch.int32, device=dev)[None, :].contiguous()
base = outside()
for T in (54, 4096):
    q = torch.randn(T, QH, D, device=dev, dtype=torch.bfloat16)
    idx = torch.stack([torch.randperm(slots, device=dev)[:TOPK].sort().values for _ in range(T)]).to(torch.int32).contiguous()
    t2r = torch.zeros(T, dtype=torch.int32, device=dev)
    fn = lambda: attend.qsa_sparse_paged_attention(q, kv[0], kv[1], idx, block_table, t2r, **extra)
    out = fn()
    ms = triton.testing.do_bench(fn, warmup=10, rep=200)
    print(f"T={T:5d} {ms * 1000:9.1f} us  checksum {out.float().abs().sum().item():.4e}", flush=True)
print(f"lmem growth {outside() - base:.0f} MiB")
for key, ck in list(orig.device_caches.values())[0][0].items():
    if ck.metadata.num_warps in (2, 4, 8):
        print(f"  compiled warps={ck.metadata.num_warps} regs={ck.n_regs} spills={ck.n_spills} shared={ck.metadata.shared}")