diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py index 3d0e893a1..7c0158ff9 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py @@ -1,9 +1,9 @@ """Step A compose support for bs>1 CP shared-KV partial-current materialize. This module holds the batch-level compose *plan* (descriptor tensors that are -identical for every layer of a batch) and the transient *arena* (tier-S carve -discipline for dense compose buffers, designed so the Step B symmetric-memory -conversion is a registration flip, not a rewrite). +identical for every layer of a batch), the rank-local dense-buffer *arena* +(Step A carve discipline), and the compact symm *staging* (Step B: the only +IPC-registered region — one round of current pages, double-buffered). Why this exists (see docs_internal/perf/cp-shared-kv-symm-materialize-design.md): the dense-buffer materialize is a partitioned gather — every byte has exactly @@ -52,10 +52,9 @@ def cp_shared_kv_compose_arena_enabled() -> bool: def cp_shared_kv_compose_symm_enabled() -> bool: - """Step B: exchange current pages via symm-heap IPC gather (zero NCCL). - - Requires the arena (the symm slab IS the arena) — checked at use site. - """ + """Step B: exchange current pages via the compact symm staging (zero + NCCL): publish my current pages to staging, barrier, gather peers'. + Independent of the arena (dense buffers stay rank-local).""" return bool(envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.get()) @@ -75,9 +74,16 @@ class ComposePlan: dummy-page convention) of all current slots in merged-span order. When the caller provides per-current-page writer (compute-owner) ranks, - ``remote_current_*`` hold the gather list for the symm exchange: only the - pages written by OTHER ranks, src page == dst page (peers' dense buffers - are at identical offsets). + the plan splits the current pages by writer for the symm staging exchange. + The staging slot of current page ``i`` is ``i`` (its index in merged-span + order) — identical on every rank, so no per-batch offset handshake: + + - ``local_current_*``: pages THIS rank wrote — published from the dense + buffer into staging slots before the barrier. The writer-ranks tensor + is carried (all entries == cp_rank) so the publish kernel can index a + uniform pointer table without a per-layer allocation. + - ``remote_current_*``: pages OTHER ranks wrote — gathered after the + barrier from ``staging[writer][slot]`` into ``dense[page]``. """ prefix_owner_ranks: torch.Tensor @@ -86,7 +92,11 @@ class ComposePlan: total_slots: int num_prefix_slots: int num_current_pages: int + local_current_writer_ranks: torch.Tensor | None = None + local_current_slot_indices: torch.Tensor | None = None + local_current_dense_pages: torch.Tensor | None = None remote_current_writer_ranks: torch.Tensor | None = None + remote_current_slot_indices: torch.Tensor | None = None remote_current_dense_pages: torch.Tensor | None = None @@ -172,7 +182,11 @@ def get_or_build_compose_plan( else: current_dense_pages = torch.empty(0, device=device, dtype=torch.long) + local_writer_ranks = None + local_slot_indices = None + local_dense_pages = None remote_writer_ranks = None + remote_slot_indices = None remote_dense_pages = None if current_page_writer_ranks is not None: num_current = int(current_dense_pages.numel()) @@ -185,8 +199,14 @@ def get_or_build_compose_plan( writers = torch.tensor( current_page_writer_ranks, dtype=torch.long, device=device ) + slot_indices = torch.arange(num_current, dtype=torch.long, device=device) remote = writers != int(layout.cp_rank) + local = ~remote + local_writer_ranks = writers[local].contiguous() + local_slot_indices = slot_indices[local].contiguous() + local_dense_pages = current_dense_pages[local].contiguous() remote_writer_ranks = writers[remote].contiguous() + remote_slot_indices = slot_indices[remote].contiguous() remote_dense_pages = current_dense_pages[remote].contiguous() plan = ComposePlan( @@ -196,7 +216,11 @@ def get_or_build_compose_plan( total_slots=total_slots, num_prefix_slots=num_prefix_slots, num_current_pages=int(current_dense_pages.numel()), + local_current_writer_ranks=local_writer_ranks, + local_current_slot_indices=local_slot_indices, + local_current_dense_pages=local_dense_pages, remote_current_writer_ranks=remote_writer_ranks, + remote_current_slot_indices=remote_slot_indices, remote_current_dense_pages=remote_dense_pages, ) plans[key] = plan @@ -208,77 +232,69 @@ def get_or_build_compose_plan( # -------------------------------------------------------------------------- -class CpComposeArena: - """Transient arena for dense compose buffers with tier-S discipline. +class _RoundParity: + """Monotonic round-epoch parity shared by the arena and the staging. - Step A behavior: a plain device slab, carved by a bump allocator that is - deterministic given the per-layer acquisition order (fixed ``kind`` order, - same shapes on every rank — shapes derive from the batch-logical slot - layout, which is rank-invariant). Buffers alternate between two halves by - layer parity so a buffer stays untouched for one extra layer. - - Step B converts this to true symmetric memory by (1) fixing ``capacity`` - at startup from the pool-derived materialize bound, (2) IPC-registering - the slab once and exchanging base pointers, (3) adding the flags page. - The carve discipline is enforced *now* so that conversion does not change - any offsets. Growth (allowed in Step A) is forbidden once registered. + A "round" is one layer-instance's compose calls (index then kv on + F-layers). Parity derives from a monotonic epoch, NOT from layer_id — + EAGLE draft layers reuse decoder layer ids, so raw-id parity would stop + alternating halves. The call sequence is identical on every rank, so + epochs (and any offsets derived from them) stay symmetric. """ - def __init__(self, device: torch.device) -> None: - self.device = device - self._slab: torch.Tensor | None = None - self._half_bytes = 0 - self._offsets = [0, 0] - self._high_water = 0 - self._registered = False # Step B: set after IPC registration - # Round/epoch state: a "round" is one layer-instance's compose calls - # (index then kv on F-layers). Parity derives from a monotonic epoch, - # NOT from layer_id — EAGLE draft layers reuse decoder layer ids, so - # raw-id parity would stop alternating halves (and same-id rounds - # would keep appending into one half across forwards). + def __init__(self) -> None: self._epoch = 0 self._round_layer_id: int | None = None self._round_kinds: set[str] = set() - # Step B symm state (populated by register_symm): - self.peer_slab_bases: torch.Tensor | None = None # CPU int64 [cp] - self.flag_ptrs: torch.Tensor | None = None # CUDA int64 [cp] - self._flags: torch.Tensor | None = None - self.cp_rank: int = -1 - - @staticmethod - def _align(nbytes: int) -> int: - return (nbytes + 255) & ~255 def begin_round(self, layer_id: int, kind: str) -> int: """Advance to a new round (and the other parity half) when needed. A new round starts when the layer id changes OR when the same buffer kind repeats for one layer id (e.g. draft layer 0 followed by the - next forward's target layer 0). Within a round, all kinds carve the - same half at deterministic offsets. The call sequence is identical - on every rank, so epochs (and therefore offsets) stay symmetric. - Returns the parity of the current round. + next forward's target layer 0). Returns the round parity. """ if layer_id != self._round_layer_id or kind in self._round_kinds: self._epoch += 1 self._round_layer_id = layer_id self._round_kinds = set() - self._offsets[self._epoch & 1] = 0 + self._on_new_round() self._round_kinds.add(kind) return self._epoch & 1 + def _on_new_round(self) -> None: + pass + + +class CpComposeArena(_RoundParity): + """Local slab for dense compose buffers (Step A allocation discipline). + + A plain device slab, carved by a bump allocator that is deterministic + given the per-layer acquisition order. Buffers alternate between two + halves by round parity so a buffer stays untouched for one extra layer. + Purely rank-local: with the compact staging exchange peers never read a + dense buffer, so the slab is never IPC-registered and may grow freely. + """ + + def __init__(self, device: torch.device) -> None: + super().__init__() + self.device = device + self._slab: torch.Tensor | None = None + self._half_bytes = 0 + self._offsets = [0, 0] + self._high_water = 0 + + @staticmethod + def _align(nbytes: int) -> int: + return (nbytes + 255) & ~255 + + def _on_new_round(self) -> None: + self._offsets[self._epoch & 1] = 0 + def _ensure_capacity(self, half_bytes: int) -> None: if self._slab is not None and half_bytes <= self._half_bytes: return - if self._registered: - raise RuntimeError( - "[CP_SHARED_KV_FAIL_FAST][compose_arena] arena growth requested " - f"after symm registration: need_half={half_bytes} " - f"have_half={self._half_bytes}. Increase " - "SGLANG_CP_SHARED_KV_SYMM_HEAP_MB (total slab MB) or check " - "the pool-derived sizing in compute_symm_half_bytes." - ) new_half = max(half_bytes, self._half_bytes * 2, 64 << 20) logger.info( "[CP-Compose-Arena] (re)allocating slab: half=%.1f MiB total=%.1f MiB", @@ -303,32 +319,65 @@ class CpComposeArena: def high_water_bytes(self) -> int: return self._high_water + +class CpComposeStaging(_RoundParity): + """Compact symm staging for the Step B current-page exchange. + + 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 holds just one round of current pages, not the whole dense buffer + (extend-bound ~100 MB instead of the pool-bound ~2.5 GB dense slab). + + Layout (fixed at registration, rank-identical, no per-batch handshake): + two parity halves, each ``[kv region | index region]``, each region + ``capacity_pages`` pages of that kind. The staging slot of current page + ``i`` is ``i`` (its index in the batch's merged current-span order, the + same on every rank), so ``peer_base + slot * page_nbytes`` addresses any + peer's copy of page ``i`` directly. + + Reuse safety is the double-buffer parity argument: my publish into half + H at round R+2 starts only after my round-R+1 barrier completed, which + proves every peer's R+1 barrier kernel ran, which (stream order) proves + every peer's round-R gather of half H finished. + """ + + KINDS = ("token_kv", "index") + + def __init__(self, device: torch.device) -> None: + super().__init__() + self.device = device + self._slab: torch.Tensor | None = None + self.capacity_pages = 0 + self.cp_size = 0 + self.cp_rank = -1 + self._page_nbytes: dict[str, int] = {} + self._region_offsets: dict[tuple[str, int], int] = {} + self.peer_slab_bases: torch.Tensor | None = None # CPU int64 [cp] + self.flag_ptrs: torch.Tensor | None = None # CUDA int64 [cp] + self._flags: torch.Tensor | None = None + + @staticmethod + def _align(nbytes: int) -> int: + return (nbytes + 255) & ~255 + @property - def symm_ready(self) -> bool: - return self._registered + def registered(self) -> bool: + return self._slab is not None - def slab_offset_of(self, buffer: torch.Tensor) -> int: - assert self._slab is not None - offset = buffer.data_ptr() - self._slab.data_ptr() - if not (0 <= offset < 2 * self._half_bytes): - raise RuntimeError( - "[CP_SHARED_KV_FAIL_FAST][compose_arena] buffer is not from " - f"this arena slab: offset={offset}" - ) - return int(offset) - - def register_symm( + def register( self, *, cp_group, cp_rank: int, cp_size: int, - half_bytes: int, + capacity_pages: int, + kv_page_nbytes: int, + index_page_nbytes: int, ) -> None: - """Fix capacity, allocate the slab + flags in CUDA-IPC memory, and + """Allocate the staging slab + barrier flags in CUDA-IPC memory and exchange handles ONCE across the CP group (collective — every rank - must call this at the same point; the first symm compose of a batch - is such a point).""" + must call this at the same point; the first symm token-KV compose of + a batch is such a point).""" from tai_kernel.nsa_prefill.ipc import ( allocate_cuda_ipc_buffer, @@ -336,17 +385,27 @@ class CpComposeArena: open_cuda_ipc_mem_handles, ) - if self._registered: + if self.registered: return - if self._slab is not None: - raise RuntimeError( - "[CP_SHARED_KV_FAIL_FAST][compose_arena] register_symm must " - "run before any non-symm slab allocation (enable " - "SGLANG_CP_SHARED_KV_COMPOSE_SYMM from startup, not mid-run)" - ) - half_bytes = self._align(half_bytes) + self.capacity_pages = int(capacity_pages) + self.cp_size = int(cp_size) + self.cp_rank = int(cp_rank) + self._page_nbytes = { + "token_kv": int(kv_page_nbytes), + "index": int(index_page_nbytes), + } + half_bytes = 0 + region_in_half = {} + for kind in self.KINDS: + region_in_half[kind] = half_bytes + half_bytes += self._align(self.capacity_pages * self._page_nbytes[kind]) + for parity in (0, 1): + for kind in self.KINDS: + self._region_offsets[(kind, parity)] = ( + parity * half_bytes + region_in_half[kind] + ) + self._slab = allocate_cuda_ipc_buffer(2 * half_bytes, device=self.device) - self._half_bytes = half_bytes flags = allocate_cuda_ipc_buffer(max(cp_size * 4, 8), device=self.device) flags.zero_() self._flags = flags @@ -369,53 +428,37 @@ class CpComposeArena: flag_peer_ptrs = _exchange(flags) torch.cuda.synchronize(self.device) self.flag_ptrs = flag_peer_ptrs.to(self.device) - self.cp_rank = int(cp_rank) - self._registered = True logger.info( - "[CP-Compose-Arena] symm slab registered: half=%.1f MiB " - "total=%.1f MiB cp_rank=%s cp_size=%s", - half_bytes / (1 << 20), + "[CP-Compose-Staging] symm staging registered: capacity=%s pages " + "(kv=%s B + index=%s B per page) total=%.1f MiB cp_rank=%s cp_size=%s", + self.capacity_pages, + self._page_nbytes["token_kv"], + self._page_nbytes["index"], 2 * half_bytes / (1 << 20), cp_rank, cp_size, ) - def peer_dense_ptrs_for(self, buffer: torch.Tensor) -> torch.Tensor: - """CPU int64 [cp]: each peer's pointer to ITS copy of ``buffer``. + def page_nbytes(self, kind: str) -> int: + return self._page_nbytes[kind] - Valid because every rank carves identical offsets (deterministic - bump over rank-invariant shapes).""" + def _offset(self, kind: str, parity: int) -> int: + return self._region_offsets[(kind, parity & 1)] + def buffer(self, kind: str, parity: int) -> torch.Tensor: + """This rank's staging region for ``kind`` at ``parity`` (uint8).""" + assert self._slab is not None + start = self._offset(kind, parity) + return self._slab[start : start + self.capacity_pages * self._page_nbytes[kind]] + + def peer_region_ptrs(self, kind: str, parity: int) -> torch.Tensor: + """CPU int64 [cp]: each peer's pointer to ITS region for this kind/parity.""" assert self.peer_slab_bases is not None - return self.peer_slab_bases + self.slab_offset_of(buffer) - - -def compute_symm_half_bytes( - *, - kv_pool_tokens: int, - page_size: int, - cp_size: int, - kv_page_nbytes: int, - index_page_nbytes: int, - slack_pages: int = 512, -) -> int: - """Pool-derived hard bound for one arena half (see design doc §1). - - Every logical page a batch can materialize is backed by the shared pool: - logical capacity = per-rank pool pages × cp_size. One half must hold the - dense KV and dense index buffers of one layer plus rounding slack. - """ - - explicit_mb = int(envs.SGLANG_CP_SHARED_KV_SYMM_HEAP_MB.get()) - if explicit_mb > 0: - return explicit_mb << 19 # MB of TOTAL slab -> half bytes - - logical_pages = (int(kv_pool_tokens) // int(page_size)) * int(cp_size) - pages = logical_pages + int(slack_pages) - return pages * (int(kv_page_nbytes) + int(index_page_nbytes)) + return self.peer_slab_bases + self._offset(kind, parity) _ARENAS: dict[torch.device, CpComposeArena] = {} +_STAGINGS: dict[torch.device, CpComposeStaging] = {} def get_compose_arena(device: torch.device) -> CpComposeArena: @@ -426,6 +469,54 @@ def get_compose_arena(device: torch.device) -> CpComposeArena: return arena +def get_compose_staging(device: torch.device) -> CpComposeStaging: + staging = _STAGINGS.get(device) + if staging is None: + staging = CpComposeStaging(device) + _STAGINGS[device] = staging + return staging + + +def compute_staging_capacity_pages( + *, + kv_pool_tokens: int, + page_size: int, + cp_size: int, + kv_page_nbytes: int, + index_page_nbytes: int, +) -> int: + """Current-page capacity of the symm staging (one parity half, per kind). + + Sized from the admission caps when available: a batch's current pages are + bounded by ceil(max_total_extend_tokens / page) plus one boundary page per + request. Falls back to the pool-derived logical-page bound (always + sufficient, just larger) when the caps are unset, e.g. in tests. + ``SGLANG_CP_SHARED_KV_SYMM_HEAP_MB`` (MB of the TOTAL slab) overrides. + """ + + explicit_mb = int(envs.SGLANG_CP_SHARED_KV_SYMM_HEAP_MB.get()) + if explicit_mb > 0: + per_page = int(kv_page_nbytes) + int(index_page_nbytes) + return max(1, (explicit_mb << 20) // (2 * per_page)) + + import sglang.srt.server_args as server_args_module + + server_args = getattr(server_args_module, "_global_server_args", None) + max_extend_tokens = ( + getattr(server_args, "cp_shared_kv_prefill_max_total_extend_tokens", None) + if server_args is not None + else None + ) + if max_extend_tokens: + max_requests = ( + getattr(server_args, "cp_shared_kv_prefill_max_batch_requests", None) + or 256 + ) + return -(-int(max_extend_tokens) // int(page_size)) + int(max_requests) + + return (int(kv_pool_tokens) // int(page_size)) * int(cp_size) + 512 + + def acquire_dense_buffer( *, device: torch.device, diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 58b0af1bc..4847c1a3c 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -11,11 +11,10 @@ import torch from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import ( acquire_dense_buffer, - compute_symm_half_bytes, - cp_shared_kv_compose_arena_enabled, + compute_staging_capacity_pages, cp_shared_kv_compose_symm_enabled, cp_shared_kv_compose_v2_enabled, - get_compose_arena, + get_compose_staging, get_or_build_compose_plan, ) from sglang.srt.layers.attention.nsa.utils import ( @@ -4481,9 +4480,7 @@ def maybe_build_current_page_writer_ranks( layer per call site would sit on the launch-critical path for nothing. """ - if not ( - cp_shared_kv_compose_symm_enabled() and cp_shared_kv_compose_arena_enabled() - ): + if not cp_shared_kv_compose_symm_enabled(): return None # The prefetchers issue their own per-span collectives and never the # barrier/symm exchange; letting them coexist would make SYMM a silent @@ -4521,38 +4518,42 @@ def maybe_build_current_page_writer_ranks( return writers -def _symm_arena_ready_or_register( +def _symm_staging_ready_or_register( *, layout: CpSharedKVLayout, kv_cache: torch.Tensor, page_size: int, ) -> bool: - """Register the symm slab on first use (collective; uniform call point). + """Register the compact symm staging on first use (collective; uniform + call point). Only the token-KV compose registers (it knows the exact KV page bytes); the index compose uses symm once registration has happened. """ - arena = get_compose_arena(kv_cache.device) - if arena.symm_ready: + staging = get_compose_staging(kv_cache.device) + if staging.registered: return True cp_group = get_attention_cp_group() # NSA index page bytes are model constants (index_head_dim 128, # quant_block 128 -> head + 4 scale bytes per token). index_page_nbytes = page_size * (128 + 4) - arena.register_symm( + kv_page_nbytes = _token_kv_page_nbytes(kv_cache, page_size) + staging.register( cp_group=cp_group, cp_rank=int(layout.cp_rank), cp_size=int(layout.cp_size), - half_bytes=compute_symm_half_bytes( + capacity_pages=compute_staging_capacity_pages( kv_pool_tokens=int(kv_cache.shape[0]), page_size=page_size, cp_size=int(layout.cp_size), - kv_page_nbytes=_token_kv_page_nbytes(kv_cache, page_size), + kv_page_nbytes=kv_page_nbytes, index_page_nbytes=index_page_nbytes, ), + kv_page_nbytes=kv_page_nbytes, + index_page_nbytes=index_page_nbytes, ) - return arena.symm_ready + return staging.registered def _symm_exchange_current_pages( @@ -4561,16 +4562,25 @@ def _symm_exchange_current_pages( layout: CpSharedKVLayout, *, page_nbytes: int, + kind: str, + layer_id: int, ) -> None: - """Step B current-page exchange: barrier + IPC gather from peers' symm - dense buffers (identical offsets on every rank). The barrier runs even - with zero remote pages — barrier COUNTS must match across ranks. + """Step B current-page exchange through the compact symm staging: + + publish my written current pages: dense[page] -> staging[slot] + barrier cp_symm_barrier (counting; release/acquire at system scope) + gather peers' pages: peer_staging[writer][slot] -> dense[page] + + The staging slot of current page ``i`` is ``i`` (merged-span order), + identical on every rank, so peers address each other's staging without a + per-batch handshake. The barrier runs even with zero local/remote pages + — barrier COUNTS must match across ranks. INVARIANT: every gate on the path to this call (compose env flags, - page_aligned metadata, symm_ready, the agreed IPC capability, prefetch - absence) MUST be rank-uniform. A per-rank divergence here desyncs the - barrier counting and hangs the CP group — never add a per-rank condition - without routing it through a group agreement first + page_aligned metadata, staging.registered, the agreed IPC capability, + prefetch absence) MUST be rank-uniform. A per-rank divergence here + desyncs the barrier counting and hangs the CP group — never add a + per-rank condition without routing it through a group agreement first (see _agreed_tai_ipc_peer_ptrs).""" from tai_kernel.nsa_prefill.ipc import ( @@ -4578,17 +4588,50 @@ def _symm_exchange_current_pages( gather_cuda_ipc_peer_pages, ) - arena = get_compose_arena(dense_buffer.device) - cp_symm_barrier(arena.flag_ptrs, self_rank=int(layout.cp_rank)) + staging = get_compose_staging(dense_buffer.device) + if int(plan.num_current_pages) > staging.capacity_pages: + # Batch-logical quantity -> raises uniformly on every rank. + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][compose_symm] batch current pages " + f"exceed the symm staging capacity: pages={plan.num_current_pages} " + f"capacity={staging.capacity_pages}. Raise " + "SGLANG_CP_SHARED_KV_SYMM_HEAP_MB or lower " + "--cp-shared-kv-prefill-max-total-extend-tokens." + ) + if page_nbytes != staging.page_nbytes(kind): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][compose_symm] page bytes diverge from " + f"the registered staging layout: kind={kind} page_nbytes=" + f"{page_nbytes} registered={staging.page_nbytes(kind)}" + ) + parity = staging.begin_round(int(layer_id), kind) if ( - plan.remote_current_dense_pages is not None - and plan.remote_current_dense_pages.numel() > 0 + plan.local_current_slot_indices is not None + and plan.local_current_slot_indices.numel() > 0 + ): + # Local copy via the gather kernel: a uniform pointer table aimed at + # my own dense buffer turns it into dense[page] -> staging[slot]. + self_ptrs = torch.full( + (staging.cp_size,), dense_buffer.data_ptr(), dtype=torch.int64 + ) + gather_cuda_ipc_peer_pages( + self_ptrs, + staging.buffer(kind, parity), + plan.local_current_writer_ranks, + plan.local_current_dense_pages, + plan.local_current_slot_indices, + page_nbytes=page_nbytes, + ) + cp_symm_barrier(staging.flag_ptrs, self_rank=int(layout.cp_rank)) + if ( + plan.remote_current_slot_indices is not None + and plan.remote_current_slot_indices.numel() > 0 ): gather_cuda_ipc_peer_pages( - arena.peer_dense_ptrs_for(dense_buffer), + staging.peer_region_ptrs(kind, parity), dense_buffer, plan.remote_current_writer_ranks, - plan.remote_current_dense_pages, + plan.remote_current_slot_indices, plan.remote_current_dense_pages, page_nbytes=page_nbytes, ) @@ -4681,7 +4724,6 @@ def _compose_token_kv_partial_current_v2( use_symm = ( cp_shared_kv_compose_symm_enabled() - and cp_shared_kv_compose_arena_enabled() and current_page_writer_ranks is not None and layer_id is not None ) @@ -4709,9 +4751,9 @@ def _compose_token_kv_partial_current_v2( ) if use_symm: # Symm requires the prefix IPC capability (same transport) and the - # registered slab; registration is collective and happens here, at - # the uniform first-use point. - use_symm = ipc_state is not None and _symm_arena_ready_or_register( + # registered staging; registration is collective and happens here, + # at the uniform first-use point. + use_symm = ipc_state is not None and _symm_staging_ready_or_register( layout=layout, kv_cache=kv_cache, page_size=page_size, @@ -4776,6 +4818,8 @@ def _compose_token_kv_partial_current_v2( plan, layout, page_nbytes=_token_kv_page_nbytes(kv_cache, page_size), + kind="token_kv", + layer_id=layer_id, ) elif plan.num_current_pages > 0: _reduce_current_pages_compact( @@ -5076,7 +5120,7 @@ def _compose_index_partial_current_v2( ) -> tuple[torch.Tensor, torch.Tensor]: """Step A/B compose for the indexer page buffer (see token-KV variant). - The index compose never registers the symm slab itself (the token-KV + The index compose never registers the symm staging itself (the token-KV compose does, knowing the exact KV page bytes); it uses symm only once registration already happened — deterministic across ranks because the layer order is identical everywhere. @@ -5084,10 +5128,9 @@ def _compose_index_partial_current_v2( use_symm = ( cp_shared_kv_compose_symm_enabled() - and cp_shared_kv_compose_arena_enabled() and current_page_writer_ranks is not None and layer_id is not None - and get_compose_arena(page_buffer.device).symm_ready + and get_compose_staging(page_buffer.device).registered ) plan = get_or_build_compose_plan( slot_remap=slot_remap, @@ -5157,6 +5200,8 @@ def _compose_index_partial_current_v2( plan, layout, page_nbytes=_page_nbytes_from_page_tensor(page_buffer), + kind="index", + layer_id=layer_id, ) elif plan.num_current_pages > 0: _reduce_current_pages_compact( diff --git a/test/manual/test_cp_shared_kv_compose_v2_8rank.py b/test/manual/test_cp_shared_kv_compose_v2_8rank.py index 3e32269bf..1c0430c5e 100644 --- a/test/manual/test_cp_shared_kv_compose_v2_8rank.py +++ b/test/manual/test_cp_shared_kv_compose_v2_8rank.py @@ -6,7 +6,7 @@ every rank with real NCCL collectives and (for v2) real tai-kernel CUDA-IPC gathers, and asserts the composed dense buffers and locs are byte-identical between the legacy per-span path and compose_v2. -Run inside the g0034 cjy-glm5-new container: +Run inside the g0033 syh-dev-new container: cd /mnt/beegfs/syh/sglang-stable && \ SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 \ PYTHONPATH=python:/mnt/beegfs/syh/tai-kernel/python \ @@ -267,38 +267,52 @@ def main() -> None: flush=True, ) - # ---- Step B: symm exchange (arena + barrier + peer gather, zero NCCL - # in the current-page phase). Multiple layers exercise parity halves. ---- + # ---- Step B: compact symm staging (publish + barrier + peer gather, + # zero NCCL in the current-page phase). Multiple layers exercise the + # staging parity halves; the arena is exercised in a second pass to + # prove the two knobs are independent. ---- + def _check_symm(layer_id: int, tag: str) -> None: + symm_kv, symm_locs = _compose( + s, layer_id=layer_id, writers=s["current_page_writer_ranks"] + ) + torch.cuda.synchronize() + assert torch.equal(ref_locs, symm_locs), ( + f"rank{rank} layer{layer_id} [{tag}]: symm locs mismatch" + ) + if not torch.equal(ref_kv, symm_kv): + diff = (ref_kv != symm_kv).any(dim=-1).any(dim=-1) + bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist() + raise AssertionError( + f"rank{rank} layer{layer_id} [{tag}]: symm dense kv mismatch " + f"at rows {bad} (of {int(diff.sum())})" + ) + with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override( True - ), envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override( - True ), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True): for layer_id in range(4): - symm_kv, symm_locs = _compose( - s, layer_id=layer_id, writers=s["current_page_writer_ranks"] - ) - torch.cuda.synchronize() - assert torch.equal(ref_locs, symm_locs), ( - f"rank{rank} layer{layer_id}: symm locs mismatch" - ) - if not torch.equal(ref_kv, symm_kv): - diff = (ref_kv != symm_kv).any(dim=-1).any(dim=-1) - bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist() - raise AssertionError( - f"rank{rank} layer{layer_id}: symm dense kv mismatch at " - f"rows {bad} (of {int(diff.sum())})" - ) + _check_symm(layer_id, "no-arena") + with envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(True): + for layer_id in range(4, 8): + _check_symm(layer_id, "arena") dist.barrier() from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import ( - get_compose_arena, + get_compose_staging, ) - assert get_compose_arena(device).symm_ready, "symm slab was not registered" + staging = get_compose_staging(device) + assert staging.registered, "symm staging was not registered" if rank == 0: + staged_mb = ( + 2 + * staging.capacity_pages + * (staging.page_nbytes("token_kv") + staging.page_nbytes("index")) + / (1 << 20) + ) print( - "PASS: symm compose byte-identical to v2 across 4 layers " - "(arena registered, barrier + peer gather engaged)", + "PASS: compact-staging symm compose byte-identical to v2 across " + f"8 layers (staging {staging.capacity_pages} pages ~{staged_mb:.1f} " + "MiB, publish + barrier + peer gather engaged; arena on and off)", flush=True, ) dist.destroy_process_group() diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_compose.py b/test/registered/unit/mem_cache/test_cp_shared_kv_compose.py index b3a7a43b2..9d0d1d483 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_compose.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_compose.py @@ -1,6 +1,6 @@ -"""Unit tests for cp_shared_kv_compose (Step A compose plan + arena). +"""Unit tests for cp_shared_kv_compose (compose plan + arena + symm staging). -Registered: CPU CI (no CUDA needed for plan/arena logic). +Registered: CPU CI (no CUDA needed for plan/arena/staging-layout logic). """ import unittest @@ -12,7 +12,9 @@ from types import SimpleNamespace from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import ( CpComposeArena, + CpComposeStaging, acquire_dense_buffer, + compute_staging_capacity_pages, get_or_build_compose_plan, ) from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -121,10 +123,15 @@ class TestComposePlanSymm(unittest.TestCase): current_page_writer_ranks=[3, 0, 3], ) # current dense pages are slots 2,3,4 -> dense 3,4,5; writers 3,0,3; - # self rank 3 -> only the page written by rank 0 is remote. + # self rank 3 -> only the page written by rank 0 is remote. The + # staging slot of current page i is i (merged-span order). self.assertEqual(plan.current_dense_pages.tolist(), [3, 4, 5]) self.assertEqual(plan.remote_current_writer_ranks.tolist(), [0]) + self.assertEqual(plan.remote_current_slot_indices.tolist(), [1]) self.assertEqual(plan.remote_current_dense_pages.tolist(), [4]) + self.assertEqual(plan.local_current_writer_ranks.tolist(), [3, 3]) + self.assertEqual(plan.local_current_slot_indices.tolist(), [0, 2]) + self.assertEqual(plan.local_current_dense_pages.tolist(), [3, 5]) def test_writer_count_mismatch_fails_fast(self): layout = CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=0) @@ -190,9 +197,10 @@ class TestComposePlanSymm(unittest.TestCase): page_size=64, layout=layout, ) + # Symm is independent of the arena: gating must work with ARENA off. with envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override( True - ), envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(True): + ), envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(False): self.assertIsNotNone( maybe_build_current_page_writer_ranks( forward_batch=fb_aligned, **kwargs @@ -251,15 +259,6 @@ class TestComposeArena(unittest.TestCase): t0 = arena.acquire(parity=p_third, nbytes=256) self.assertEqual(t0.data_ptr(), d0.data_ptr()) # back to half A - def test_growth_is_forbidden_after_registration(self): - arena = CpComposeArena(torch.device("cpu")) - parity = arena.begin_round(0, "kv") - arena.acquire(parity=parity, nbytes=64) - arena._registered = True - parity = arena.begin_round(1, "kv") - with self.assertRaises(RuntimeError): - arena.acquire(parity=parity, nbytes=arena._half_bytes + 1) - def test_acquire_dense_buffer_plain_alloc_when_arena_disabled(self): with envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(False): buf = acquire_dense_buffer( @@ -286,5 +285,90 @@ class TestComposeArena(unittest.TestCase): self.assertTrue(buf.is_contiguous()) +class TestComposeStaging(unittest.TestCase): + """Layout/parity logic of the compact symm staging (no IPC needed: the + region offsets and round parity are pure functions of the registration + parameters and the call sequence).""" + + def _staging(self, capacity_pages=4, kv_nbytes=1000, index_nbytes=300): + staging = CpComposeStaging(torch.device("cpu")) + # Bypass the IPC allocation: install the layout exactly as register() + # computes it, backed by a plain CPU slab. + staging.capacity_pages = capacity_pages + staging.cp_size = 8 + staging.cp_rank = 0 + staging._page_nbytes = {"token_kv": kv_nbytes, "index": index_nbytes} + half = 0 + region_in_half = {} + for kind in CpComposeStaging.KINDS: + region_in_half[kind] = half + half += staging._align(capacity_pages * staging._page_nbytes[kind]) + for parity in (0, 1): + for kind in CpComposeStaging.KINDS: + staging._region_offsets[(kind, parity)] = ( + parity * half + region_in_half[kind] + ) + staging._slab = torch.zeros(2 * half, dtype=torch.uint8) + staging.peer_slab_bases = torch.full((8,), 1 << 20, dtype=torch.int64) + return staging, half + + def test_regions_are_disjoint_and_parity_halves_do_not_overlap(self): + staging, half = self._staging() + kv0 = staging.buffer("token_kv", 0) + idx0 = staging.buffer("index", 0) + kv1 = staging.buffer("token_kv", 1) + # kv and index regions of one half are adjacent but disjoint + # (index starts at the 256-aligned end of kv). + self.assertEqual(idx0.data_ptr() - kv0.data_ptr(), staging._align(4 * 1000)) + # The parity-1 half starts exactly one half past parity 0. + self.assertEqual(kv1.data_ptr() - kv0.data_ptr(), half) + self.assertEqual(kv0.numel(), 4 * 1000) + self.assertEqual(idx0.numel(), 4 * 300) + + def test_peer_region_ptrs_offset_matches_local_layout(self): + staging, half = self._staging() + base = staging.peer_slab_bases + self.assertTrue( + torch.equal(staging.peer_region_ptrs("token_kv", 0), base) + ) + self.assertTrue( + torch.equal( + staging.peer_region_ptrs("index", 1), + base + half + staging._align(4 * 1000), + ) + ) + + def test_round_parity_alternates_and_shares_round_across_kinds(self): + staging, _ = self._staging() + p0 = staging.begin_round(0, "index") + self.assertEqual(staging.begin_round(0, "token_kv"), p0) + p1 = staging.begin_round(1, "index") + self.assertNotEqual(p1, p0) + # Repeated (id, kind) -> new round (EAGLE draft id reuse). + p_again = staging.begin_round(1, "index") + self.assertNotEqual(p_again, p1) + + def test_capacity_pages_from_env_override_caps_and_pool(self): + kwargs = dict( + kv_pool_tokens=64 * 1000, + page_size=64, + cp_size=8, + kv_page_nbytes=41_984, + index_page_nbytes=8_448, + ) + with envs.SGLANG_CP_SHARED_KV_SYMM_HEAP_MB.override(101): + pages = compute_staging_capacity_pages(**kwargs) + self.assertEqual( + pages, (101 << 20) // (2 * (41_984 + 8_448)) + ) + # No caps set in unit tests -> pool-derived fallback. + with envs.SGLANG_CP_SHARED_KV_SYMM_HEAP_MB.override(0): + import sglang.srt.server_args as server_args_module + + if getattr(server_args_module, "_global_server_args", None) is None: + pages = compute_staging_capacity_pages(**kwargs) + self.assertEqual(pages, 1000 * 8 + 512) + + if __name__ == "__main__": unittest.main()