diff --git a/test/manual/bench_mqa_logits_chunking.py b/test/manual/bench_mqa_logits_chunking.py new file mode 100644 index 000000000..80ab5d406 --- /dev/null +++ b/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()