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