Restructure bs>1 CP shared-KV compose to one gather + one collective

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>
This commit is contained in:
2026-06-11 19:16:28 +00:00
co-authored by Claude Fable 5
parent ecf7998c80
commit f47b739e30
6 changed files with 1357 additions and 4 deletions
+9
View File
@@ -220,6 +220,15 @@ class Envs:
# large bs) but coarser overlap. 1 = per-layer.
SGLANG_CP_SHARED_KV_PER_LAYER_GROUP = EnvInt(8)
SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
# Step A compose restructure for bs>1 partial-current materialize: one IPC
# prefix gather + one compact-current collective per buffer per layer
# (fast path), or one whole-buffer all-reduce (fallback), instead of one
# all-reduce per request span. 0 restores the legacy per-span collectives.
SGLANG_CP_SHARED_KV_COMPOSE_V2 = EnvBool(True)
# Carve compose dense buffers from the tier-S arena (deterministic bump,
# layer-parity halves). Default off in Step A; Step B registers the arena
# as symmetric memory and flips this on.
SGLANG_CP_SHARED_KV_COMPOSE_ARENA = EnvBool(False)
# NSA MQA logits are materialized as fp32 [q, k] buffers inside DeepGEMM.
# Lower values split query rows more aggressively to cap peak temporary memory.
SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION = EnvFloat(0.2)