Cut per-layer CPU on the prefill launch path: validators, plans, spans

From the nsys CPU-gap attribution (launch thread, one 78-layer forward: 374ms API time; 642 cudaStreamSynchronize blocking 89.5ms and overlapping 122ms of the 505ms GPU idle; ~44ms pure-Python before concat_mla_absorb_q): - memory_pool_host: skip validate_page_aligned_token_indices on CUDA tensors in _get_indexer_page_indices and _prepare_load_page_indices — torch.any/torch.equal there cost a queue-deep cudaStreamSynchronize per layer-group submit (~0.42ms each, ~12.7ms/forward measured). Same construction-based-invariant guard the CacheController pair check already documents; CPU/test tensors stay validated. - nsa_indexer: per-batch _CpRaggedIndexPlan replaces the per-F-layer rebuild of the O(total-q-tokens) topk offset list and the 6-7 int32 ragged descriptor tensors (segment records, kv_lens/q_starts/q_lens/ k_bases/q_bases/current_bases). All inputs are batch metadata; the plan is anchored on the forward batch with a content key over cp_index. - nsa_indexer forward_indexer: read seq_lens_cpu instead of a device seq_lens[i].item() per request per layer (one stream sync each). - cp_shared_kv_runtime: get_or_build_batch_slot_spans caches the layer-invariant prefix/current slot spans per batch (the builders read logical_pages only for its shape); nsa_backend x3 + nsa_indexer call sites switched. Microbenchmark (idle H200, traced batch shape bs=12 / 44.6K q tokens, test/manual/bench_cpu_gap_fixes.py, equality-checked): validator path 197.1us -> 59.2us per submit under a busy queue (x3.3); ragged plan 3238.6us -> 36.2us per layer (x90, ~128ms launch-thread time per forward at 40 F-layers); slot spans 20.1us -> 0.5us (x41). Layer suites A/B vs HEAD: identical failure set (5 pre-existing CPU-tensor indexer tests), no regressions. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-11 21:01:00 +00:00
parent b356773d2f
commit a24111a5f4
5 changed files with 519 additions and 159 deletions
@@ -0,0 +1,220 @@
+#!/usr/bin/env python3
+"""Micro-benchmark for the pre-attention CPU-gap fixes (task #14).
+
+Scenario mirrors the traced production batch: bs=12, prefix 640..26304 tok,
+extends ~3.7K tok, cp8 in-seq-split (2*cp segments/request), page 64.
+
+  1. page-aligned validator skip on CUDA tensors
+     (memory_pool_host._get_indexer_page_indices hot path) — measured with a
+     busy GPU queue, because torch.any/.equal sync for the whole queue.
+  2. ragged index descriptor plan: rebuild-per-layer (old) vs per-batch cache.
+  3. slot-span builders: rebuild-per-layer (old) vs per-batch cache.
+
+Run (single GPU is enough):
+  PYTHONPATH=python python test/manual/bench_cpu_gap_fixes.py
+"""
+from __future__ import annotations
+
+import time
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.layers.attention.nsa import nsa_indexer
+from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
+    build_batch_current_slot_spans,
+    build_batch_prefix_slot_spans,
+    get_or_build_batch_slot_spans,
+)
+from sglang.srt.mem_cache.page_index_utils import (
+    validate_page_aligned_token_indices,
+)
+
+DEV = torch.device("cuda", 0)
+PAGE = 64
+BS = 12
+CP = 8
+PREFIX_LENS = [19200, 256] + [26304] * 10
+EXTEND_LENS = [3776, 7360] + [3347] * 10
+REPS = 200
+
+
+def timed(fn, reps=REPS, warmup=20):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(reps):
+        fn()
+    torch.cuda.synchronize()
+    return (time.perf_counter() - t0) / reps * 1e6  # us
+
+
+def make_fake_batch():
+    fb = SimpleNamespace()
+    fb.seq_lens_cpu = torch.tensor(
+        [p + e for p, e in zip(PREFIX_LENS, EXTEND_LENS)], dtype=torch.int64
+    )
+    fb.extend_seq_lens_cpu = list(EXTEND_LENS)
+    fb.extend_prefix_lens_cpu = list(PREFIX_LENS)
+    return fb
+
+
+def make_cp_index():
+    # 2*CP zigzag segments per request over the extend, page-aligned-ish.
+    cp_index = []
+    for req, extend in enumerate(EXTEND_LENS):
+        seg = max(PAGE, (extend // (2 * CP)) // PAGE * PAGE)
+        pos = 0
+        while pos < extend:
+            end = min(pos + seg, extend)
+            cp_index.append((req, pos, end))
+            pos = end
+    return cp_index
+
+
+def bench_validator():
+    n_pages = 600  # ~one layer-group submit worth of pages
+    starts = torch.arange(n_pages, device=DEV, dtype=torch.int64) * PAGE
+    indices = (
+        starts[:, None] + torch.arange(PAGE, device=DEV, dtype=torch.int64)
+    ).reshape(-1)
+
+    # Busy queue: enqueue ~0.5ms of GEMM before each validator call, the way
+    # the real submit lands behind a layer's compute.
+    a = torch.randn(2048, 2048, device=DEV, dtype=torch.bfloat16)
+    b = torch.randn(2048, 2048, device=DEV, dtype=torch.bfloat16)
+
+    def old_path():
+        for _ in range(4):
+            a @ b
+        validate_page_aligned_token_indices(indices, PAGE, "bench")
+        starts2 = indices.reshape(-1, PAGE)[:, 0] // PAGE
+        return starts2
+
+    def new_path():
+        for _ in range(4):
+            a @ b
+        if not indices.is_cuda:
+            validate_page_aligned_token_indices(indices, PAGE, "bench")
+        starts2 = indices.reshape(-1, PAGE)[:, 0] // PAGE
+        return starts2
+
+    # measure WALL time per call without trailing torch.cuda.synchronize in
+    # the loop (the sync inside the validator is exactly what we measure).
+    def wall(fn, reps=60, warmup=10):
+        for _ in range(warmup):
+            fn()
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        for _ in range(reps):
+            fn()
+        t1 = time.perf_counter()  # NO sync: CPU-side blocking is the metric
+        torch.cuda.synchronize()
+        return (t1 - t0) / reps * 1e6
+
+    old = wall(old_path)
+    new = wall(new_path)
+    print(
+        f"1. validator (busy queue, {n_pages} pages): old={old:8.1f}us "
+        f"new={new:8.1f}us  speedup x{old/new:.1f}"
+    )
+
+
+def bench_ragged_plan():
+    fb = make_fake_batch()
+    cp_index = make_cp_index()
+
+    def old_build():
+        # the pre-fix per-layer work: full python loop + 7 tensor H2Ds
+        return nsa_indexer._build_cp_ragged_index_plan(fb, cp_index, DEV, None)
+
+    fb2 = make_fake_batch()
+
+    def cached():
+        return nsa_indexer._get_or_build_cp_ragged_index_plan(
+            fb2, cp_index, DEV, None
+        )
+
+    old = timed(old_build)
+    new = timed(cached)
+    n_tokens = sum(EXTEND_LENS)
+    print(
+        f"2. ragged index plan (bs={BS}, {len(cp_index)} segs, {n_tokens} q tok): "
+        f"per-layer rebuild={old:8.1f}us cached={new:8.1f}us  speedup x{old/new:.0f}"
+    )
+
+
+def bench_spans():
+    fb = make_fake_batch()
+    pages_per_req = max(
+        (p + e + PAGE - 1) // PAGE for p, e in zip(PREFIX_LENS, EXTEND_LENS)
+    )
+    logical_pages = torch.zeros((BS, pages_per_req), dtype=torch.int64)
+
+    def old_build():
+        prefix = build_batch_prefix_slot_spans(
+            logical_pages=logical_pages,
+            prefix_lens_cpu=PREFIX_LENS,
+            page_size=PAGE,
+        )
+        current = build_batch_current_slot_spans(
+            logical_pages=logical_pages,
+            prefix_lens_cpu=PREFIX_LENS,
+            extend_lens_cpu=EXTEND_LENS,
+            page_size=PAGE,
+        )
+        return prefix, current
+
+    def cached():
+        return get_or_build_batch_slot_spans(
+            fb,
+            logical_pages=logical_pages,
+            prefix_lens_cpu=PREFIX_LENS,
+            extend_lens_cpu=EXTEND_LENS,
+            page_size=PAGE,
+            want_prefix=True,
+        )
+
+    old = timed(old_build, reps=2000)
+    new = timed(cached, reps=2000)
+    print(
+        f"3. slot spans (bs={BS}): per-layer rebuild={old:8.1f}us "
+        f"cached={new:8.1f}us  speedup x{old/new:.0f}"
+    )
+
+
+def main():
+    torch.cuda.init()
+    print(f"device: {torch.cuda.get_device_name(0)}")
+    bench_validator()
+    bench_ragged_plan()
+    bench_spans()
+    # equality check: cached plan tensors match a fresh build
+    fb = make_fake_batch()
+    cp_index = make_cp_index()
+    p1 = nsa_indexer._build_cp_ragged_index_plan(fb, cp_index, DEV, None)
+    p2 = nsa_indexer._get_or_build_cp_ragged_index_plan(fb, cp_index, DEV, None)
+    assert torch.equal(p1.topk_indices_offset_override, p2.topk_indices_offset_override)
+    assert torch.equal(p1.kv_lens, p2.kv_lens) and torch.equal(p1.q_bases, p2.q_bases)
+    assert p1.segment_records == p2.segment_records
+    s1 = build_batch_current_slot_spans(
+        logical_pages=torch.zeros((BS, 512), dtype=torch.int64),
+        prefix_lens_cpu=PREFIX_LENS,
+        extend_lens_cpu=EXTEND_LENS,
+        page_size=PAGE,
+    )
+    _, s2 = get_or_build_batch_slot_spans(
+        SimpleNamespace(),
+        logical_pages=torch.zeros((BS, 512), dtype=torch.int64),
+        prefix_lens_cpu=PREFIX_LENS,
+        extend_lens_cpu=EXTEND_LENS,
+        page_size=PAGE,
+        want_prefix=False,
+    )
+    assert s1 == s2
+    print("EQUALITY CHECKS PASS")
+
+
+if __name__ == "__main__":
+    main()