Seed the IPC agreement before prefetch hit/miss can diverge

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 <noreply@anthropic.com>

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
 
 

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,

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