The publish-variant staging exchange measured ~equal to the default
compact-current AR (90.4 vs 86.5 ms/batch on the traced scenario): the
publish copy and the barrier serialized behind the 0.65 ms prefix
gather ate the transport win the isolated current exchange showed
(0.196 vs 0.354 ms). Fix the structure instead of the copy: current
rows are now written straight INTO the staging — the fill kernels take
their write destinations solely from page_inverse, so a per-batch
staging-remapped page inverse on the plan retargets them with zero
kernel changes — then cp_symm_barrier, then ONE slot-dense gather
covers prefix pages (pool pointers) and ALL current pages (staging
pointers, including this rank's own) through a concatenated 2*cp
pointer table where current slots carry owner = cp_size + writer and
src = staging slot. No publish copy, no prefix pre-gather, no second
gather.
The fused fill's loc outputs are dense-geometry-bound, so the token-KV
path computes mixed_locs/staging row indices once per batch (they are
layer-invariant) and the per-layer fill collapses to a single
index_copy_ into the zeroed staging span.
Benchmark (g0033 8xH200, byte-exact, idle-checked): 62.8 ms/batch vs
84.4 default Step A (-26%) and 60.8 ideal; publish variant was 88.0.
151 unit tests; 8-rank GPU byte-exactness vs v2 across 8 layers,
arena on and off.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Peers only ever read the CURRENT pages of a rank's compose output — the
prefix comes straight from the IPC-registered KV pool — so the symm
region does not need to hold the whole dense buffer (pool-bound ~2.5 GB
double-buffered slab). It now holds one round of current pages in
merged-span order (extend-cap-bound, ~58-100 MB), and dense buffers
become purely rank-local (plain allocations or the optional local arena;
COMPOSE_SYMM no longer requires COMPOSE_ARENA).
Exchange per compose call: publish my written current pages
dense[page] -> staging[slot i] (slot = the page's batch current index,
identical on every rank, so peers address each other's staging with no
per-batch handshake), cp_symm_barrier, gather peers'
staging[writer][slot] -> dense[page] via the existing src!=dst page
gather. Reuse safety keeps the parity-half distance-2 argument, now on
the staging. Capacity sizing comes from the admission caps
(max_total_extend_tokens / max_batch_requests) with a pool-derived
fallback and the SYMM_HEAP_MB override; overflow fails fast
(batch-logical, hence rank-uniform).
Idea credit: laoyao0822's touched-pages-proportional staging
(906ecbe5d4), rebound onto our barrier-gated, group-agreed transport.
Validated on g0033 8xH200: 151 unit tests; 8-rank GPU byte-exactness
vs compose_v2 across 8 layers (arena on and off, parity halves
exercised); benchmark path e (real protocol) byte-exact, current-page
exchange 0.196 ms vs 0.354 ms compact-AR isolated.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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>
Replaces the last remaining NCCL collective in the bs>1 compose layer loop
(the compact-current all-reduce) with a DeepEP-style counting barrier plus
a CUDA-IPC gather from peers' symmetric dense buffers, behind
SGLANG_CP_SHARED_KV_COMPOSE_SYMM (+ COMPOSE_ARENA, both default off).
- CpComposeArena.register_symm: fixes capacity at the pool-derived bound
(logical pages x dense page unit, overridable via
SGLANG_CP_SHARED_KV_SYMM_HEAP_MB), allocates slab + flags in CUDA-IPC
memory, exchanges handles once over the CP group; deterministic bump
carve means peer_base + my_offset addresses any peer's dense buffer with
no per-layer handshakes. Registration happens only in the token-KV
compose (uniform first-use point); growth after registration raises.
- Current pages are single-writer at page granularity under the
page-aligned in-seq split, so the exchange is the existing
gather_cuda_ipc_peer_pages with src==dst page ids and writer (compute
owner) ranks; writers are built per batch by
build_batch_current_page_writer_ranks and gated by
maybe_build_current_page_writer_ranks (page_aligned metadata required).
The barrier runs even with zero remote pages (counts must match).
- _agreed_tai_ipc_peer_ptrs: the per-rank IPC capability probe is now
agreed across the CP group (one-time MIN all-reduce per pool tensor) so
ranks can never split between gather and collective paths and deadlock
on mismatched NCCL shapes.
- ComposePlan cache re-anchored ON the slot_remap object (forward-batch
lifetime) instead of a module-level tensor-identity key, which could go
stale when a freed tensor's address is reused by the next batch.
Validation (g0034 syh-dev-new): tai-kernel cp_symm_barrier correctness
(200 adversarial iterations, rotating 10ms producer delays, phase-safety,
flags drained at quiescence) and perf (7.7us max-rank latency) both pass;
8-rank GPU test extends to the symm path - byte-identical to compose_v2
across 4 layers (parity halves exercised, slab registered); mem_cache
suite 464 passed with the only failures being a documented pre-existing
sys.modules stub pollution pair, reproduced identically with
SGLANG_CP_SHARED_KV_COMPOSE_V2=0.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The bs>1 partial-current materialize issued one sum-all-reduce per request
span per buffer per layer (48 collectives per F-layer at bs=12, 2880 per
batch, 419ms + 49ms launch gaps in the production trace), all inline on
the compute stream. The data is a partitioned gather, not a reduction:
every byte has exactly one producer.
compose_v2 (SGLANG_CP_SHARED_KV_COMPOSE_V2, default on) replaces this with:
- fast path: one tai-kernel CUDA-IPC slot-dense gather covering ALL prefix
spans (full-range descriptors, -1 sentinels zero-fill current slots and
replace the dense zero-fill) + ONE collective over the compact current
pages (uint8 byte view; exact because every byte is writer-exclusive).
- fallback (no peer IPC): local materialize of all prefix spans + ONE
whole-buffer sum-all-reduce (rows are still writer-exclusive pre-reduce).
IPC capability is decided once by the cached peer-pointer probe; after a
successful probe a failing gather raises (no per-call try/except).
cp_shared_kv_compose.py adds the per-batch ComposePlan descriptor cache
(layer-invariant, keyed on the slot_logical_pages identity) and the
CpComposeArena with tier-S carve discipline (deterministic bump, layer-
parity halves; default off) so the Step B symmetric-memory conversion is
a registration flip.
Microbenchmark (g0034 8xH200, traced 12-req batch, per batch): per-span
214ms -> fused AR 119ms -> IPC prefix + compact current 84ms; symm target
61ms. Validation: 143 unit tests incl. v2-contract twins, legacy siblings
and rank-merged simulations under both paths; mem_cache dir 432 passed;
8-rank GPU byte-exactness vs legacy with real NCCL + IPC (test/manual/
test_cp_shared_kv_compose_v2_8rank.py) passed with no fallback markers.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
