Add mqa-logits chunking cost micro-benchmark (专题 S2a)

Measures the real cost of shrinking SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB:
the faithful indexer loop (deep_gemm.fp8_mqa_logits +
fast_topk_transform_fused, serial) at GLM-5.1 shapes (H=32, D=128,
topk=2048) across cold-chunk / tail-chunk / warm-continuation /
warm-long scenarios. g0033 1xH200 results: 2GB costs at most +5.5%
(cold 64K chunk) and is -6.7% on the heaviest warm-long shape; the
knee is ~1GB; 0.5GB is +30%. Shrinking 8->2GB frees ~6GB of the
per-batch CP admission budget for KV layer buffers.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

test/manual/bench_mqa_logits_chunking.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""Micro-benchmark: cost of chunking the NSA indexer mqa-logits/topk loop.
+
+专题 S2(a) (docs_internal/perf/prefill-compute-intensity-plan.md): the 18G CP
+admission budget is dominated by the SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB=8
+worst-case term.  Before shrinking it, measure how much per-layer indexer
+time the smaller chunks actually cost — the loop is
+``deep_gemm.fp8_mqa_logits`` (logits buffer = rows x kv x 4B) followed by
+``sgl_kernel.fast_topk_transform_fused`` per chunk
+(nsa_indexer._mqa_logits_topk_ragged_chunked), serial on one stream.
+
+GLM-5.1 shapes: index_n_heads=32, index_head_dim=128, index_topk=2048,
+attn-cp8 in-seq split (per-rank q rows = chunk/8; the rank sees the FULL
+dense kv).  Single GPU — the loop is rank-local.
+
+Run (g0033 syh-dev-new, 1 GPU):
+  cd /mnt/beegfs/syh/sglang-stable && \
+  PYTHONPATH=python:/mnt/beegfs/syh/tai-kernel/python \
+  python test/manual/bench_mqa_logits_chunking.py \
+  2>&1 | tee /mnt/beegfs/syh/log/bench_mqa_chunking_$(date +%Y%m%d_%H%M%S).log
+"""
+from __future__ import annotations
+
+import time
+
+import deep_gemm
+import torch
+from sgl_kernel import fast_topk_transform_fused
+
+H = 32  # index_n_heads
+D = 128  # index_head_dim
+TOPK = 2048
+BYTES_PER_ELEM = 4  # fp32 logits
+
+# (name, prefix_tokens, chunk_extend_tokens) — global lens; per-rank q rows =
+# extend/8 (cp8 in-seq split, this rank takes the LAST slice = worst-case ke).
+SCENARIOS = [
+    ("cold_chunk1_64K", 0, 65536),
+    ("80K_tail_chunk_14K", 65536, 14464),
+    ("warm_cont_150K+4K", 150_000, 4096),
+    ("warm_long_150K+60K", 150_000, 60000),
+]
+CHUNK_GBS = [0.0, 8.0, 4.0, 2.0, 1.0, 0.5]  # 0.0 = one-shot reference
+
+
+def _build(prefix: int, extend: int, device):
+    n = prefix + extend  # full dense kv this rank sees
+    m = max(1, extend // 8)  # per-rank q rows
+    rank_start = prefix + (extend // 8) * 7  # last slice global offset
+    q = torch.randn(m, H, D, device=device).clamp(-2, 2).to(torch.float8_e4m3fn)
+    kv = torch.randn(n, D, device=device).clamp(-2, 2).to(torch.float8_e4m3fn)
+    kv_scales = torch.rand(n, device=device) + 0.5
+    weights = torch.rand(m, H, device=device)
+    ks = torch.zeros(m, dtype=torch.int32, device=device)
+    ke = (
+        torch.arange(m, dtype=torch.int32, device=device) + rank_start + 1
+    ).clamp(max=n)
+    # token-granular page table for one request (slot i = token i) + paged
+    # topk inputs, mirroring the PAGED topk_transform path.
+    page_table_1 = torch.arange(n, dtype=torch.int32, device=device).view(1, n)
+    cu_seqlens_q = torch.tensor([0, m], dtype=torch.int32, device=device)
+    return dict(
+        m=m, n=n, q=q, kv=kv, kv_scales=kv_scales, weights=weights,
+        ks=ks, ke=ke, page_table_1=page_table_1, cu_seqlens_q=cu_seqlens_q,
+    )
+
+
+def _run_chunked(s, max_rows: int) -> torch.Tensor:
+    m = s["m"]
+    out = None
+    start = 0
+    while start < m:
+        end = min(start + max_rows, m)
+        logits = deep_gemm.fp8_mqa_logits(
+            s["q"][start:end],
+            (s["kv"], s["kv_scales"]),
+            s["weights"][start:end],
+            s["ks"][start:end],
+            s["ke"][start:end],
+            clean_logits=False,
+        )
+        cu = torch.tensor(
+            [0, end - start], dtype=torch.int32, device=logits.device
+        )
+        topk = fast_topk_transform_fused(
+            score=logits,
+            lengths=s["ke"][start:end],
+            page_table_size_1=s["page_table_1"],
+            cu_seqlens_q=cu,
+            topk=TOPK,
+            row_starts=s["ks"][start:end],
+        )
+        if out is None:
+            out = topk.new_full((m, topk.shape[1]), -1)
+        out[start:end] = topk
+        start = end
+    return out
+
+
+def main() -> None:
+    device = torch.device("cuda", 0)
+    torch.cuda.set_device(device)
+    print(f"device={torch.cuda.get_device_name(device)} H={H} D={D} topk={TOPK}")
+    for name, prefix, extend in SCENARIOS:
+        s = _build(prefix, extend, device)
+        m, n = s["m"], s["n"]
+        print(
+            f"\n=== {name}: prefix={prefix} extend={extend} "
+            f"per-rank q_rows={m} kv={n} ==="
+        )
+        ref_ms = None
+        for gb in CHUNK_GBS:
+            if gb <= 0:
+                max_rows = m
+            else:
+                max_rows = max(1, int(gb * 1e9) // (n * BYTES_PER_ELEM))
+                max_rows = min(max_rows, m)
+            num_chunks = -(-m // max_rows)
+            peak_mb = max_rows * n * BYTES_PER_ELEM / 1e6
+            # warmup
+            for _ in range(3):
+                _run_chunked(s, max_rows)
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            reps = 10
+            for _ in range(reps):
+                _run_chunked(s, max_rows)
+            torch.cuda.synchronize()
+            ms = (time.perf_counter() - t0) / reps * 1000
+            if ref_ms is None:
+                ref_ms = ms
+            label = "one-shot" if gb <= 0 else f"{gb:>4.1f}GB"
+            print(
+                f"  {label}: chunks={num_chunks:>3} rows/chunk={max_rows:>6} "
+                f"peak_logits={peak_mb:>8.1f}MB  {ms:8.3f} ms/layer "
+                f"({(ms / ref_ms - 1) * 100:+6.1f}% vs one-shot)"
+            )
+
+
+if __name__ == "__main__":
+    main()