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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user