leavelet a24111a5f4 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

6.8 KiB

Raw Blame History

View Raw