From 9c6cf278b03194b0db6ef066c4178e3c32db809b Mon Sep 17 00:00:00 2001 From: leavelet Date: Thu, 11 Jun 2026 22:56:31 +0000 Subject: [PATCH] Seed the IPC agreement before prefetch hit/miss can diverge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Opus review of the symm prefetch wiring found a latent group deadlock: the prefetch hit path never touched _agreed_tai_ipc_peer_ptrs, so the first use of a layer's pool tensor — which issues a capability MIN all-reduce plus an IPC handle all-gather — only happened on the sync fallback. With every rank hitting from batch one (the prefetcher runs a layer ahead), a later lone-rank miss would be the only rank issuing those collectives. Both consumes now seed the (idempotent, per-pool- tensor-cached) agreement at entry, before any per-rank early return; the index consume gets the pool buffer from its indexer call site. Unreachable in today's config (prefetch env off) but a blocker for ever enabling it. Also per review: the [pool|staging] combined pointer-table cache now keys per pool table instead of holding one — the token and index pool tables alternate every layer and evicted each other, so the cache never hit; and a debug-gated check that current locs all map into the staging page inverse (index_copy_ has no skip semantics for stray -1). Validated on g0033: 151 unit tests; all four 8-rank byte-exactness phases. Co-Authored-By: Claude Fable 5 --- .../attention/nsa/cp_shared_kv_compose.py | 34 ++++++++++--------- .../attention/nsa/cp_shared_kv_prefetch.py | 31 +++++++++++++++++ .../srt/layers/attention/nsa/nsa_indexer.py | 1 + 3 files changed, 50 insertions(+), 16 deletions(-) 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 931b7e625..05cd1138f 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 @@ -389,10 +389,15 @@ class CpComposeStaging(_RoundParity): self._flags: torch.Tensor | None = None # Launch-path caches: the per-(kind, parity) peer pointer tables and # the [pool | staging] concatenations are tiny CPU tensors, but - # rebuilding them per layer per kind is avoidable overhead. + # rebuilding them per layer per kind is avoidable overhead. The + # combined cache keys by pool-table id AND stores the table reference + # in the value: the ref pins the identity (no stale-address reuse), + # and per-table entries mean the token/index pool tables alternating + # every layer do not evict each other. self._region_peer_ptrs: dict[tuple[str, int], torch.Tensor] = {} - self._combined_pool_table: torch.Tensor | None = None - self._combined_ptrs: dict[tuple[str, int], torch.Tensor] = {} + self._combined_ptrs: dict[ + tuple[int, str, int], tuple[torch.Tensor, torch.Tensor] + ] = {} @staticmethod def _align(nbytes: int) -> int: @@ -506,21 +511,18 @@ class CpComposeStaging(_RoundParity): ) -> torch.Tensor: """CPU int64 [2*cp]: ``[pool peer ptrs | staging peer ptrs]``. - Cached per (kind, parity) against the (long-lived, per-pool) agreed - pointer table — holding the table reference pins its identity, so - the identity check cannot go stale on address reuse. + Cached per (pool table, kind, parity); the agreed pool pointer + tables are long-lived (one per pool tensor) so the cache stays tiny. """ - key = (kind, parity & 1) - if self._combined_pool_table is not pool_peer_ptrs: - self._combined_pool_table = pool_peer_ptrs - self._combined_ptrs = {} - table = self._combined_ptrs.get(key) - if table is None: - table = torch.cat( - [pool_peer_ptrs, self.peer_region_ptrs(kind, parity)] - ).contiguous() - self._combined_ptrs[key] = table + key = (id(pool_peer_ptrs), kind, parity & 1) + entry = self._combined_ptrs.get(key) + if entry is not None and entry[0] is pool_peer_ptrs: + return entry[1] + table = torch.cat( + [pool_peer_ptrs, self.peer_region_ptrs(kind, parity)] + ).contiguous() + self._combined_ptrs[key] = (pool_peer_ptrs, table) return table diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index 0cf56dbf0..51207448a 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -13,6 +13,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import ( get_or_build_compose_plan, ) from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + _agreed_tai_ipc_peer_ptrs, _all_reduce_materialized_buffer_async, _all_reduce_materialized_buffer_range, _page_nbytes_from_page_tensor, @@ -78,6 +79,24 @@ class _PrefetchSymmState: setattr(self, name, kwargs.get(name)) +def _seed_symm_ipc_agreement(prefetcher: Any, pool_buffer: torch.Tensor) -> None: + """Seed the group-agreed IPC capability for this layer's pool tensor + BEFORE any per-rank hit/miss divergence. + + The sync-compose fallback's first use of a pool tensor issues collectives + (capability MIN all-reduce + IPC handle all-gather); if only a lone miss + rank ran them the CP group would deadlock. Seeding here is uniform — the + consume call itself is batch-logical — and idempotent (cached per pool + tensor).""" + + if ( + prefetcher.symm_writers is not None + and prefetcher.layout.cp_size > 1 + and cp_shared_kv_compose_symm_enabled() + ): + _agreed_tai_ipc_peer_ptrs(pool_buffer, prefetcher.layout) + + def _prefetch_symm_active(prefetcher: Any, device: torch.device) -> bool: """Rank-uniform gate for the prefetch-path symm current exchange. @@ -680,6 +699,14 @@ class CpSharedKVMlaPrefetcher: page_size=self.page_size, loc_req_id=current_req_id.reshape(-1), ).to(torch.long) + if cp_shared_kv_debug_enabled() and staging_current_rows.numel() > 0: + # index_copy_ has no skip semantics for stray -1s (debug-only: + # the min() syncs). + if int(staging_current_rows.min().item()) < 0: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][prefetch_symm] current locs " + "map outside the staging page inverse" + ) # mixed_locs is layer-invariant; the fused fill that used to produce # it builds masks sized by its target buffer, which is now the # staging — so compute it once here in logical space instead. @@ -891,6 +918,7 @@ class CpSharedKVMlaPrefetcher: in ``logical_locs`` are remapped to the appended current KV rows. """ + _seed_symm_ipc_agreement(self, kv_cache) if self.disabled: self._log_layer( layer_id, @@ -1840,7 +1868,10 @@ class CpSharedKVIndexPrefetcher: page_size: int, index_head_dim: int, current_req_id: torch.Tensor | None = None, + pool_page_buffer: torch.Tensor | None = None, ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + if pool_page_buffer is not None: + _seed_symm_ipc_agreement(self, pool_page_buffer) if self.disabled: self._log_layer( layer_id, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 835d8d544..6c06b9cb7 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -794,6 +794,7 @@ class Indexer(MultiPlatformOp): page_size=page_size, index_head_dim=forward_batch.token_to_kv_pool.index_head_dim, current_req_id=current_req_id, + pool_page_buffer=index_buffer, ) if prefetched is not None: return prefetched