#!/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_pipelined(s, max_rows: int, side: torch.cuda.Stream) -> torch.Tensor: """Same loop with topk_N on a side stream overlapping logits_{N+1}. Settles whether the small-chunk penalty is serialization (pipelining helps) or small-M GEMM inefficiency (it cannot) — both kernels are SM-heavy, so concurrent streams may just timeshare. """ m = s["m"] main = torch.cuda.current_stream() out = None bounds = [] start = 0 while start < m: end = min(start + max_rows, m) bounds.append((start, end)) start = end def _logits(b): return deep_gemm.fp8_mqa_logits( s["q"][b[0] : b[1]], (s["kv"], s["kv_scales"]), s["weights"][b[0] : b[1]], s["ks"][b[0] : b[1]], s["ke"][b[0] : b[1]], clean_logits=False, ) def _topk(logits, b): cu = torch.tensor([0, b[1] - b[0]], dtype=torch.int32, device=logits.device) return fast_topk_transform_fused( score=logits, lengths=s["ke"][b[0] : b[1]], page_table_size_1=s["page_table_1"], cu_seqlens_q=cu, topk=TOPK, row_starts=s["ks"][b[0] : b[1]], ) prev_logits, prev_b, prev_evt = None, None, None for b in bounds: if prev_logits is not None: # Launch topk(N-1) on the side stream gated ONLY on its own # logits event, then immediately launch logits(N) on main — the # two overlap (or timeshare SMs; that is what we are measuring). side.wait_event(prev_evt) with torch.cuda.stream(side): topk = _topk(prev_logits, prev_b) if out is None: out = topk.new_full((m, topk.shape[1]), -1) out[prev_b[0] : prev_b[1]] = topk prev_logits.record_stream(side) logits = _logits(b) evt = torch.cuda.Event() evt.record(main) prev_logits, prev_b, prev_evt = logits, b, evt side.wait_event(prev_evt) with torch.cuda.stream(side): topk = _topk(prev_logits, prev_b) if out is None: out = topk.new_full((m, topk.shape[1]), -1) out[prev_b[0] : prev_b[1]] = topk prev_logits.record_stream(side) main.wait_stream(side) out.record_stream(main) return out 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 pipe_ms = None if num_chunks > 1: side = torch.cuda.Stream() # NOTE: the kernel chain is run-to-run nondeterministic even # serial-vs-serial (near-equal fp32 selection), so byte # equality is meaningless here; sanity-check shape only and # leave correctness to the indexer's own validation paths. ref = _run_chunked(s, max_rows) got = _run_chunked_pipelined(s, max_rows, side) torch.cuda.synchronize() assert ref.shape == got.shape for _ in range(3): _run_chunked_pipelined(s, max_rows, side) torch.cuda.synchronize() t0 = time.perf_counter() for _ in range(reps): _run_chunked_pipelined(s, max_rows, side) torch.cuda.synchronize() pipe_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" pipe_str = ( f" pipelined {pipe_ms:8.3f} ms ({(pipe_ms / ms - 1) * 100:+6.1f}%)" if pipe_ms is not None else "" ) 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)" + pipe_str ) if __name__ == "__main__": main()