Add symm-heap current-page exchange for CP shared-KV compose (Step B)

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>
This commit is contained in:
2026-06-11 20:13:04 +00:00
co-authored by Claude Fable 5
parent f47b739e30
commit c55e406176
8 changed files with 626 additions and 68 deletions
+7
View File
@@ -229,6 +229,13 @@ class Envs:
# 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)
# Step B: exchange current pages via barrier + symm-heap IPC gather (zero
# NCCL in the layer loop). Requires COMPOSE_ARENA=1, page-aligned in-seq
# split, and tai-kernel cp_symm_barrier. Default off until e2e parity.
SGLANG_CP_SHARED_KV_COMPOSE_SYMM = EnvBool(False)
# Total symm slab size override in MB (both halves). 0 = pool-derived
# hard bound (logical pages x dense page unit, see design doc).
SGLANG_CP_SHARED_KV_SYMM_HEAP_MB = EnvInt(0)
# 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)