Wire the symm current exchange into the CP shared-KV prefetchers

The bs=1 MLA/index prefetchers replaced their consume-side trailing-range
NCCL all-reduce with the staging exchange: fill current rows straight into
this round's staging span (token-KV collapses to one cached index_copy;
the index fill kernel is just pointed at the staging page inverse),
cp_symm_barrier, then gather ALL current pages — this rank's own included
— from the stagings into the prefetched dense buffer.  The symm+prefetcher
FAIL_FAST is gone.

Rank-uniformity moves with it: staging registration now also happens in
maybe_create (batch-logical gates, before any per-rank miss can diverge),
because with a prefetcher active the sync compose runs only on per-rank
misses and its lazy collective registration would hang.  A hit/miss
divergence itself stays barrier-safe — both the prefetch consume and the
sync-compose fallback execute exactly one begin_round + barrier per
(layer, kind), and the counting barrier is shape-free (unlike the AR pair
it replaces, which would shape-mismatch).

Found by the new index test phase: the fill/remap kernel family skips
page id 0 as the SGLang dummy page, so a 0-based first staging slot was
never written.  The staging layout now reserves row 0 (slot of current
page i = i + 1) for every kind, matching the convention instead of
depending on per-kernel behavior.

Launch-path cost: per-(kind,parity) peer pointer tables and the
[pool|staging] concatenations are precomputed/cached (identity pinned by
holding the pool-table reference); all prefetch descriptors, staging row
indices, and mixed_locs are built once per batch.

Validated on g0033 8xH200: 151 unit tests; 8-rank byte-exactness for
token sync symm (8 layers), index sync symm (4 layers, new phase), and
MLA + index prefetch consume_prefix_with_current vs the legacy sync
compose.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/layers/attention/nsa/cp_shared_kv_compose.py

@@ -76,14 +76,14 @@ class ComposePlan:
 
     When the caller provides per-current-page writer (compute-owner) ranks,
     the plan additionally carries the symm staging descriptors.  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:
+    slot of current page ``i`` is ``i + 1`` (merged-span order; row 0 = dummy)
+    — identical on every rank, so no per-batch offset handshake:
 
     - ``symm_all_owner_ranks`` / ``symm_all_src_pages``: ONE slot-dense
       gather descriptor pair covering the whole dense buffer, indexed against
       a concatenated ``[pool peer ptrs | staging peer ptrs]`` table — prefix
       slots keep the pool owner/page, current slots carry
-      ``cp_size + writer`` / staging slot ``i``, everything else stays -1
+      ``cp_size + writer`` / staging slot ``i + 1``, everything else stays -1
       (sentinel zero-fill).
     - ``staging_page_inverse``: ``page_inverse`` with current pages remapped
       to their staging slot (others -1), so the unchanged fill kernels write
@@ -115,11 +115,13 @@ def _build_staging_page_inverse(
 
     ``page_inverse`` maps (request row, logical page) -> dense page id; the
     fill kernels derive every current-row write destination from it.  In the
-    staging copy a current dense page becomes its index in the (ascending)
-    merged-span order — the staging slot — and every other entry becomes -1,
-    so the unchanged fill kernels write current rows into the compact staging
-    viewed as a page-major buffer (no dummy page).  Current locs only ever
-    map to current pages, so the -1 entries are never exercised by the fill.
+    staging copy a current dense page becomes ONE PLUS its index in the
+    (ascending) merged-span order — the staging slot — and every other entry
+    becomes -1.  The +1 keeps staging row 0 unused: the fill/remap kernel
+    family treats page id 0 as the SGLang dummy page and silently skips
+    writes to it, so a 0-based first staging slot would never be filled.
+    Current locs only ever map to current pages, so the -1 entries are never
+    exercised by the fill.
     """
 
     if page_inverse.dim() == 1:
@@ -132,7 +134,7 @@ def _build_staging_page_inverse(
         current_dense_pages, page_inverse.clamp(min=0).contiguous()
     ).clamp(max=num_current - 1)
     is_current = (page_inverse > 0) & (current_dense_pages[pos] == page_inverse)
-    return torch.where(is_current, pos, invalid)
+    return torch.where(is_current, pos + 1, invalid)
 
 
 def get_or_build_compose_plan(
@@ -230,8 +232,10 @@ def get_or_build_compose_plan(
             )
             current_slots = current_dense_pages - 1
             symm_all_owner_ranks[current_slots] = writers + int(layout.cp_size)
+            # Staging slot of current page i is i + 1 (row 0 = dummy; see
+            # _build_staging_page_inverse).
             symm_all_src_pages[current_slots] = torch.arange(
-                num_current, dtype=torch.long, device=device
+                1, num_current + 1, dtype=torch.long, device=device
             )
         page_inverse = getattr(slot_remap, "page_inverse", None)
         if page_inverse is not None:
@@ -357,10 +361,11 @@ class CpComposeStaging(_RoundParity):
 
     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.
+    ``capacity_pages + 1`` pages of that kind (row 0 unused — the fill/remap
+    kernel family skips page id 0 as the dummy page).  The staging slot of
+    current page ``i`` is ``i + 1`` (one plus 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 fill into half
     H at round R+2 starts only after my round-R+1 barrier completed, which
@@ -382,6 +387,12 @@ class CpComposeStaging(_RoundParity):
         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
+        # 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.
+        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] = {}
 
     @staticmethod
     def _align(nbytes: int) -> int:
@@ -425,7 +436,10 @@ class CpComposeStaging(_RoundParity):
         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])
+            # +1: staging row 0 stays unused (dummy-page convention).
+            half_bytes += self._align(
+                (self.capacity_pages + 1) * self._page_nbytes[kind]
+            )
         for parity in (0, 1):
             for kind in self.KINDS:
                 self._region_offsets[(kind, parity)] = (
@@ -455,6 +469,10 @@ class CpComposeStaging(_RoundParity):
         flag_peer_ptrs = _exchange(flags)
         torch.cuda.synchronize(self.device)
         self.flag_ptrs = flag_peer_ptrs.to(self.device)
+        for key, offset in self._region_offsets.items():
+            self._region_peer_ptrs[key] = (
+                self.peer_slab_bases + offset
+            ).contiguous()
         logger.info(
             "[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",
@@ -476,12 +494,34 @@ class CpComposeStaging(_RoundParity):
         """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]]
+        nbytes = (self.capacity_pages + 1) * self._page_nbytes[kind]
+        return self._slab[start : start + nbytes]
 
     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._offset(kind, parity)
+        return self._region_peer_ptrs[(kind, parity & 1)]
+
+    def combined_ptr_table(
+        self, pool_peer_ptrs: torch.Tensor, kind: str, parity: int
+    ) -> 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.
+        """
+
+        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
+        return table
 
 
 _ARENAS: dict[torch.device, CpComposeArena] = {}

python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py

@@ -7,9 +7,20 @@ from typing import Any, Optional
 
 import torch
 
+from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
+    cp_shared_kv_compose_symm_enabled,
+    get_compose_staging,
+    get_or_build_compose_plan,
+)
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     _all_reduce_materialized_buffer_async,
     _all_reduce_materialized_buffer_range,
+    _page_nbytes_from_page_tensor,
+    _symm_begin_current_staging,
+    _symm_staging_ready_or_register,
+    _token_kv_page_nbytes,
+    build_current_loc_remap,
+    build_current_page_mask,
     cp_shared_kv_debug_enabled,
     cp_shared_kv_mla_prefetch_enabled,
     cp_shared_kv_mla_prefetch_log,
@@ -26,6 +37,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     get_or_build_shared_token_kv_slot_remap,
     materialize_local_paged_buffer_page_slots_into,
     materialize_local_token_kv_page_slots_into,
+    maybe_build_current_page_writer_ranks,
     remap_logical_pages_to_slot_dense_pages,
     remap_logical_locs_to_slot_dense_locs_optimized,
     slot_range_to_page_slice,
@@ -38,6 +50,54 @@ from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
 logger = logging.getLogger(__name__)
 
 
+class _PrefetchSymmState:
+    """Layer-invariant symm-exchange descriptors for one prefetch batch.
+
+    ``num_current_pages`` / ``staging_page_inverse`` make this object a valid
+    ``plan`` argument for ``_symm_begin_current_staging``.  The gather covers
+    ALL current pages (this rank's own come from its own staging — the fill
+    wrote them there, not into the dense buffer), so unlike the sync mega
+    gather no prefix descriptors are involved: the prefix was already
+    materialized by the prefetch stream.
+    """
+
+    __slots__ = (
+        "num_current_pages",
+        "staging_page_inverse",
+        "writer_ranks",
+        "staging_slots",
+        "current_dense_pages",
+        "page_nbytes",
+        "staging_current_rows",
+        "mixed_locs",
+        "dense_pages",
+    )
+
+    def __init__(self, **kwargs) -> None:
+        for name in self.__slots__:
+            setattr(self, name, kwargs.get(name))
+
+
+def _prefetch_symm_active(prefetcher: Any, device: torch.device) -> bool:
+    """Rank-uniform gate for the prefetch-path symm current exchange.
+
+    Every condition is batch-logical or group-agreed: ``symm_writers`` comes
+    from ``maybe_build_current_page_writer_ranks`` (env + batch metadata),
+    and ``registered`` only flips inside a collective registration.  A
+    prefetch hit/miss divergence across ranks stays barrier-safe because the
+    sync-compose fallback also runs exactly one begin_round + barrier per
+    (layer, kind).
+    """
+
+    return (
+        prefetcher.symm_writers is not None
+        and prefetcher.layout.cp_size > 1
+        and prefetcher.prefix_pages < prefetcher.total_slots
+        and cp_shared_kv_compose_symm_enabled()
+        and get_compose_staging(device).registered
+    )
+
+
 def _prefetch_log(message: str, *args) -> None:
     cp_shared_kv_mla_prefetch_log(message, *args)
 
@@ -335,6 +395,8 @@ class CpSharedKVMlaPrefetcher:
         owned_prefix_pages: int = -1,
         owned_total_pages: int = -1,
         stream: Optional[torch.cuda.Stream] = None,
+        slot_remap: Any = None,
+        symm_writers: Optional[list] = None,
     ) -> None:
         self.layout = layout
         self.page_size = page_size
@@ -352,6 +414,9 @@ class CpSharedKVMlaPrefetcher:
         self.pending_attention_handle: Optional[CpSharedKVMlaPrefetchHandle] = None
         self.disabled = False
         self._cpu_timing = _PrefetchCpuTiming()
+        self.slot_remap = slot_remap
+        self.symm_writers = symm_writers
+        self._symm_state: Optional[_PrefetchSymmState] = None
 
     @classmethod
     def maybe_create(
@@ -509,6 +574,25 @@ class CpSharedKVMlaPrefetcher:
             logger.exception("Failed to initialize CP shared KV MLA prefetcher.")
             return None
 
+        symm_writers = None
+        if cp_shared_kv_compose_symm_enabled() and layout.cp_size > 1:
+            symm_writers = maybe_build_current_page_writer_ranks(
+                forward_batch=forward_batch,
+                prefix_lens_cpu=extend_prefix_lens_cpu,
+                extend_lens_cpu=extend_seq_lens_cpu,
+                page_size=page_size,
+                layout=layout,
+            )
+            if symm_writers is not None:
+                # Collective registration at a batch-uniform point: with a
+                # prefetcher active the sync compose only runs on per-rank
+                # misses, so its lazy first-compose registration would
+                # diverge.  Must NOT be swallowed — a half-registered group
+                # is a hang, not a fallback.
+                _symm_staging_ready_or_register(
+                    layout=layout, kv_cache=kv_cache, page_size=page_size
+                )
+
         owned_prefix_pages = _debug_owned_pages_count(
             layout, remap.slot_logical_pages[:prefix_pages]
         )
@@ -556,8 +640,82 @@ class CpSharedKVMlaPrefetcher:
             owned_prefix_pages=owned_prefix_pages,
             owned_total_pages=owned_total_pages,
             stream=prefetch_stream,
+            slot_remap=remap,
+            symm_writers=symm_writers,
         )
 
+    def _get_or_build_symm_state(
+        self,
+        *,
+        kv_cache: torch.Tensor,
+        logical_locs: torch.Tensor,
+        current_locs: torch.Tensor,
+        loc_req_id: torch.Tensor,
+        current_req_id: torch.Tensor,
+    ) -> _PrefetchSymmState:
+        state = self._symm_state
+        if state is not None:
+            return state
+        plan = get_or_build_compose_plan(
+            slot_remap=self.slot_remap,
+            layout=self.layout,
+            physical_page_capacity=kv_cache.shape[0] // self.page_size,
+            prefix_spans=[(0, self.prefix_pages)],
+            current_spans=[(self.prefix_pages, self.total_slots)],
+            kind="token_kv",
+            current_page_writer_ranks=self.symm_writers,
+        )
+        num_current = int(plan.num_current_pages)
+        writer_ranks = (
+            plan.symm_all_owner_ranks.index_select(0, plan.current_dense_pages - 1)
+            - int(self.layout.cp_size)
+        ).contiguous()
+        # Staging slot of current page i is i + 1 (row 0 = dummy page).
+        staging_slots = torch.arange(
+            1, num_current + 1, dtype=torch.long, device=writer_ranks.device
+        )
+        staging_current_rows = remap_logical_locs_to_slot_dense_locs_optimized(
+            current_locs.reshape(-1),
+            page_inverse=plan.staging_page_inverse,
+            page_size=self.page_size,
+            loc_req_id=current_req_id.reshape(-1),
+        ).to(torch.long)
+        # 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.
+        logical_locs = filter_locs_mappable_to_physical_pool(
+            logical_locs=logical_locs,
+            layout=self.layout,
+            physical_token_capacity=kv_cache.shape[0],
+        )
+        dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
+            logical_locs,
+            page_inverse=self.page_inverse,
+            page_size=self.page_size,
+            loc_req_id=loc_req_id,
+        )
+        current_mask, _ = build_current_loc_remap(logical_locs, current_locs)
+        current_page_mask = build_current_page_mask(
+            logical_locs, current_locs, page_size=self.page_size
+        )
+        mixed_locs = torch.where(
+            current_page_mask & (~current_mask),
+            torch.full_like(dense_locs, -1),
+            dense_locs,
+        )
+        state = _PrefetchSymmState(
+            num_current_pages=num_current,
+            staging_page_inverse=plan.staging_page_inverse,
+            writer_ranks=writer_ranks,
+            staging_slots=staging_slots,
+            current_dense_pages=plan.current_dense_pages,
+            page_nbytes=_token_kv_page_nbytes(kv_cache, self.page_size),
+            staging_current_rows=staging_current_rows,
+            mixed_locs=mixed_locs,
+        )
+        self._symm_state = state
+        return state
+
     def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool:
         start_layer = int(getattr(token_to_kv_pool, "start_layer", 0))
         kv_buffer = getattr(token_to_kv_pool, "kv_buffer", None)
@@ -782,15 +940,81 @@ class CpSharedKVMlaPrefetcher:
         dense_kv_cache = handle.dense_kv_cache
 
         remap_cpu = _cpu_timing_start()
+        if loc_req_id is None:
+            loc_req_id = torch.zeros_like(logical_locs, dtype=torch.long)
+        if current_req_id is None:
+            current_req_id = torch.zeros_like(current_locs, dtype=torch.long)
+
+        if _prefetch_symm_active(self, dense_kv_cache.device):
+            # Symm current exchange: fill current rows straight into this
+            # round's staging span, barrier, gather ALL current pages from
+            # the stagings (this rank's own included) into the prefetched
+            # dense buffer.  Zero NCCL; descriptors and mixed_locs are
+            # layer-invariant and cached on the prefetcher.
+            from tai_kernel.nsa_prefill.ipc import (
+                cp_symm_barrier,
+                gather_cuda_ipc_peer_pages,
+            )
+
+            state = self._get_or_build_symm_state(
+                kv_cache=kv_cache,
+                logical_locs=logical_locs,
+                current_locs=current_locs,
+                loc_req_id=loc_req_id,
+                current_req_id=current_req_id,
+            )
+            staging = get_compose_staging(dense_kv_cache.device)
+            parity, staging_span = _symm_begin_current_staging(
+                staging=staging,
+                plan=state,
+                kind="token_kv",
+                layer_id=layer_id,
+                page_nbytes=state.page_nbytes,
+            )
+            num_rows = int(state.staging_current_rows.numel())
+            if num_rows > 0:
+                staging_rows = staging_span.view(kv_cache.dtype).reshape(
+                    (state.num_current_pages + 1) * self.page_size,
+                    *kv_cache.shape[1:],
+                )
+                staging_rows.index_copy_(
+                    0,
+                    state.staging_current_rows,
+                    current_kv_cache[:num_rows],
+                )
+            cp_symm_barrier(
+                staging.flag_ptrs, self_rank=int(self.layout.cp_rank)
+            )
+            gather_cuda_ipc_peer_pages(
+                staging.peer_region_ptrs("token_kv", parity),
+                dense_kv_cache,
+                state.writer_ranks,
+                state.staging_slots,
+                state.current_dense_pages,
+                page_nbytes=state.page_nbytes,
+            )
+            remap_ms = _cpu_timing_ms(remap_cpu)
+            total_ms = _cpu_timing_ms(consume_cpu)
+            self._log_layer(
+                layer_id,
+                "consume_prefix_current_hit layer=%s prefix_pages=%s "
+                "dense_rows=%s current_rows=%s symm=1 total_ms=%.3f "
+                "wait_ms=%.3f remap_ms=%.3f",
+                layer_id,
+                self.prefix_pages,
+                int(dense_kv_cache.shape[0]),
+                int(current_kv_cache.shape[0]),
+                total_ms,
+                wait_ms,
+                remap_ms,
+            )
+            return dense_kv_cache, state.mixed_locs
+
         logical_locs = filter_locs_mappable_to_physical_pool(
             logical_locs=logical_locs,
             layout=self.layout,
             physical_token_capacity=kv_cache.shape[0],
         )
-        if loc_req_id is None:
-            loc_req_id = torch.zeros_like(logical_locs, dtype=torch.long)
-        if current_req_id is None:
-            current_req_id = torch.zeros_like(current_locs, dtype=torch.long)
         dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
             logical_locs,
             page_inverse=self.page_inverse,
@@ -1120,6 +1344,8 @@ class CpSharedKVIndexPrefetcher:
         owned_prefix_pages: int = -1,
         owned_total_pages: int = -1,
         stream: Optional[torch.cuda.Stream] = None,
+        slot_remap: Any = None,
+        symm_writers: Optional[list] = None,
     ) -> None:
         self.layout = layout
         self.prefix_pages = prefix_pages
@@ -1136,6 +1362,9 @@ class CpSharedKVIndexPrefetcher:
         self.pending_attention_handle: Optional[CpSharedKVIndexPrefetchHandle] = None
         self.disabled = False
         self._cpu_timing = _PrefetchCpuTiming()
+        self.slot_remap = slot_remap
+        self.symm_writers = symm_writers
+        self._symm_state: Optional[_PrefetchSymmState] = None
 
     @classmethod
     def maybe_create(
@@ -1335,6 +1564,31 @@ class CpSharedKVIndexPrefetcher:
             logger.exception("Failed to initialize CP shared KV index prefetcher.")
             return None
 
+        symm_writers = None
+        if cp_shared_kv_compose_symm_enabled() and layout.cp_size > 1:
+            symm_writers = maybe_build_current_page_writer_ranks(
+                forward_batch=forward_batch,
+                prefix_lens_cpu=extend_prefix_lens_cpu,
+                extend_lens_cpu=extend_seq_lens_cpu,
+                page_size=page_size,
+                layout=layout,
+            )
+            if symm_writers is not None:
+                # Registration sizing needs the token-KV page bytes, so fetch
+                # the key buffer; a no-op when the MLA prefetcher (created
+                # first) already registered.  Collective — must not be
+                # swallowed (see the MLA twin).
+                _symm_staging_ready_or_register(
+                    layout=layout,
+                    kv_cache=_prefetch_pool_get_key_buffer(
+                        token_to_kv_pool=token_to_kv_pool,
+                        layer_id=first_layer_id,
+                        stream=prefetch_stream,
+                        path="index_symm_register",
+                    ),
+                    page_size=page_size,
+                )
+
         owned_prefix_pages = _debug_owned_pages_count(
             layout, remap.slot_logical_pages[:prefix_pages]
         )
@@ -1378,8 +1632,55 @@ class CpSharedKVIndexPrefetcher:
             owned_prefix_pages=owned_prefix_pages,
             owned_total_pages=owned_total_pages,
             stream=prefetch_stream,
+            slot_remap=remap,
+            symm_writers=symm_writers,
         )
 
+    def _get_or_build_symm_state(
+        self,
+        *,
+        dense_page_buffer: torch.Tensor,
+        logical_pages: torch.Tensor,
+    ) -> _PrefetchSymmState:
+        state = self._symm_state
+        if state is not None:
+            return state
+        plan = get_or_build_compose_plan(
+            slot_remap=self.slot_remap,
+            layout=self.layout,
+            physical_page_capacity=None,
+            prefix_spans=[(0, self.prefix_pages)],
+            current_spans=[(self.prefix_pages, self.total_slots)],
+            kind="index",
+            current_page_writer_ranks=self.symm_writers,
+        )
+        num_current = int(plan.num_current_pages)
+        writer_ranks = (
+            plan.symm_all_owner_ranks.index_select(0, plan.current_dense_pages - 1)
+            - int(self.layout.cp_size)
+        ).contiguous()
+        # Staging slot of current page i is i + 1 (row 0 = dummy page).
+        staging_slots = torch.arange(
+            1, num_current + 1, dtype=torch.long, device=writer_ranks.device
+        )
+        # The returned dense-pages remap is layer-invariant too.
+        dense_pages = remap_logical_pages_to_slot_dense_pages(
+            logical_pages,
+            page_inverse=self.page_inverse,
+            page_req_id=build_page_table_row_req_id(logical_pages),
+        )
+        state = _PrefetchSymmState(
+            num_current_pages=num_current,
+            staging_page_inverse=plan.staging_page_inverse,
+            writer_ranks=writer_ranks,
+            staging_slots=staging_slots,
+            current_dense_pages=plan.current_dense_pages,
+            page_nbytes=_page_nbytes_from_page_tensor(dense_page_buffer),
+            dense_pages=dense_pages,
+        )
+        self._symm_state = state
+        return state
+
     def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool:
         start_layer = int(getattr(token_to_kv_pool, "start_layer", 0))
         kv_buffer = getattr(token_to_kv_pool, "kv_buffer", None)
@@ -1599,6 +1900,72 @@ class CpSharedKVIndexPrefetcher:
         remap_cpu = _cpu_timing_start()
         if current_req_id is None:
             current_req_id = torch.zeros_like(current_locs, dtype=torch.long)
+
+        if _prefetch_symm_active(self, dense_page_buffer.device):
+            # Symm current exchange (see the MLA twin): fill into the
+            # staging, barrier, gather all current pages into the prefetched
+            # dense page buffer.  Zero NCCL.
+            from tai_kernel.nsa_prefill.ipc import (
+                cp_symm_barrier,
+                gather_cuda_ipc_peer_pages,
+            )
+
+            state = self._get_or_build_symm_state(
+                dense_page_buffer=dense_page_buffer,
+                logical_pages=logical_pages,
+            )
+            staging = get_compose_staging(dense_page_buffer.device)
+            parity, staging_span = _symm_begin_current_staging(
+                staging=staging,
+                plan=state,
+                kind="index",
+                layer_id=layer_id,
+                page_nbytes=state.page_nbytes,
+            )
+            if state.num_current_pages > 0:
+                staging_pages = staging_span.view(
+                    dense_page_buffer.dtype
+                ).reshape(
+                    state.num_current_pages + 1, *dense_page_buffer.shape[1:]
+                )
+                fill_current_index_page_slots(
+                    dense_page_buffer=staging_pages,
+                    current_index_k=current_index_k,
+                    current_index_scale=current_index_scale,
+                    current_locs=current_locs,
+                    page_inverse=state.staging_page_inverse,
+                    page_size=page_size,
+                    index_head_dim=index_head_dim,
+                    current_req_id=current_req_id,
+                )
+            cp_symm_barrier(
+                staging.flag_ptrs, self_rank=int(self.layout.cp_rank)
+            )
+            gather_cuda_ipc_peer_pages(
+                staging.peer_region_ptrs("index", parity),
+                dense_page_buffer,
+                state.writer_ranks,
+                state.staging_slots,
+                state.current_dense_pages,
+                page_nbytes=state.page_nbytes,
+            )
+            remap_ms = _cpu_timing_ms(remap_cpu)
+            total_ms = _cpu_timing_ms(consume_cpu)
+            self._log_layer(
+                layer_id,
+                "index_consume_prefix_current_hit layer=%s prefix_pages=%s "
+                "dense_pages=%s current_rows=%s symm=1 total_ms=%.3f "
+                "wait_ms=%.3f remap_ms=%.3f",
+                layer_id,
+                self.prefix_pages,
+                int(dense_page_buffer.shape[0]),
+                int(current_index_k.shape[0]),
+                total_ms,
+                wait_ms,
+                remap_ms,
+            )
+            return dense_page_buffer, state.dense_pages
+
         dense_pages = remap_logical_pages_to_slot_dense_pages(
             logical_pages,
             page_inverse=self.page_inverse,

python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py

@@ -4482,22 +4482,6 @@ def maybe_build_current_page_writer_ranks(
 
     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
-    # no-op on every prefetch hit. Fail fast instead of measuring parity.
-    if (
-        getattr(forward_batch, "cp_shared_kv_mla_prefetcher", None) is not None
-        or getattr(forward_batch, "cp_shared_kv_index_prefetcher", None) is not None
-        or envs.SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH.get()
-    ):
-        raise RuntimeError(
-            "[CP_SHARED_KV_FAIL_FAST][compose_symm] "
-            "SGLANG_CP_SHARED_KV_COMPOSE_SYMM requires the CP shared-KV "
-            "prefetchers to be disabled (the prefetch path bypasses the "
-            "symm exchange and would keep issuing per-span all-reduces). "
-            "Unset SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH / index prefetch "
-            "or disable COMPOSE_SYMM."
-        )
     metadata = getattr(forward_batch, "nsa_cp_metadata", None)
     if metadata is None or not getattr(metadata, "page_aligned", False):
         return None
@@ -4659,8 +4643,10 @@ def _symm_begin_current_staging(
             "staging page inverse (slot_remap.page_inverse missing?)"
         )
     parity = staging.begin_round(int(layer_id), kind)
+    # +1: staging row 0 is the (unused) dummy page; zeroing it keeps any
+    # accidental slot-0 read deterministic.
     span = staging.buffer(kind, parity)[
-        : int(plan.num_current_pages) * page_nbytes
+        : (int(plan.num_current_pages) + 1) * page_nbytes
     ]
     span.zero_()
     return parity, span
@@ -4701,11 +4687,8 @@ def _symm_barrier_and_gather_all(
     from tai_kernel.nsa_prefill.ipc import cp_symm_barrier
 
     cp_symm_barrier(staging.flag_ptrs, self_rank=int(layout.cp_rank))
-    combined_ptrs = torch.cat(
-        [pool_peer_ptrs, staging.peer_region_ptrs(kind, parity)]
-    )
     kernels.materialize_cuda_ipc_peer_pages_slot_dense(
-        combined_ptrs,
+        staging.combined_ptr_table(pool_peer_ptrs, kind, parity),
         dense_buffer,
         plan.symm_all_owner_ranks,
         plan.symm_all_src_pages,
@@ -4894,7 +4877,8 @@ def _compose_token_kv_partial_current_v2(
         num_rows = int(fill_state.staging_current_rows.numel())
         if num_rows > 0:
             staging_rows = staging_span.view(kv_cache.dtype).reshape(
-                int(plan.num_current_pages) * page_size, *kv_cache.shape[1:]
+                (int(plan.num_current_pages) + 1) * page_size,
+                *kv_cache.shape[1:],
             )
             staging_rows.index_copy_(
                 0,
@@ -5332,7 +5316,7 @@ def _compose_index_partial_current_v2(
         )
         if plan.num_current_pages > 0:
             staging_pages = staging_span.view(page_buffer.dtype).reshape(
-                int(plan.num_current_pages), *page_buffer.shape[1:]
+                int(plan.num_current_pages) + 1, *page_buffer.shape[1:]
             )
             fill_current_index_page_slots(
                 dense_page_buffer=staging_pages,

test/manual/test_cp_shared_kv_compose_v2_8rank.py

@@ -183,6 +183,7 @@ def _build_scenario(rank: int, cp_size: int, device: torch.device):
 
     return dict(
         kv_cache=kv_cache,
+        logical_pages=logical_pages,
         logical_locs=logical_locs.reshape(-1),
         loc_req_id=loc_req_id,
         current_kv=current_kv,
@@ -197,6 +198,325 @@ def _build_scenario(rank: int, cp_size: int, device: torch.device):
     )
 
 
+def _build_index_scenario(rank: int, cp_size: int, device: torch.device, s):
+    """Index page buffer + current K/scale rows over the bs=4 geometry."""
+    page_size = s["page_size"]
+    index_page_nbytes = page_size * (128 + 4)
+    physical_pages = int(s["kv_cache"].shape[0]) // page_size
+    page_buffer = torch.zeros(
+        (physical_pages, index_page_nbytes), dtype=torch.uint8, device=device
+    )
+    # Owned PREFIX pages get deterministic payloads (mirrors the KV pool).
+    slot_pages = s["slot_remap"].slot_logical_pages.reshape(-1).cpu()
+    prefix_slots = set()
+    for start, end in s["prefix_slot_spans"]:
+        prefix_slots.update(range(int(start), int(end)))
+    for slot in sorted(prefix_slots):
+        logical_page = int(slot_pages[slot])
+        if logical_page <= 0 or (logical_page - 1) % cp_size != rank:
+            continue
+        phys_page = (logical_page - 1) // cp_size + 1
+        page_buffer[phys_page] = (
+            (torch.arange(index_page_nbytes, dtype=torch.int64) + logical_page * 37)
+            .remainder_(251)
+            .to(torch.uint8)
+            .to(device)
+        )
+    rows = int(s["current_locs"].numel())
+    current_k = (
+        (
+            torch.arange(rows, dtype=torch.int64).view(-1, 1) * 13
+            + torch.arange(128, dtype=torch.int64).view(1, -1)
+            + rank * 17
+        )
+        .remainder_(249)
+        .to(torch.uint8)
+        .to(device)
+    )
+    current_scale = (
+        torch.arange(rows, dtype=torch.float32, device=device).view(-1, 1) * 0.25
+        + 1.0
+        + rank
+    )
+    from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as rt
+
+    slot_remap = rt.build_shared_paged_buffer_slot_remap(
+        page_buffer,
+        s["logical_pages"],
+        s["layout"],
+    )
+    return dict(
+        page_buffer=page_buffer,
+        current_k=current_k,
+        current_scale=current_scale,
+        current_locs=s["current_locs"],
+        current_req_id=s["current_req_id"],
+        slot_remap=slot_remap,
+    )
+
+
+def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None:
+    """bs=1 prefetch-path symm consume vs the legacy sync compose (MLA+index)."""
+    from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as rt
+    from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import (
+        CpSharedKVIndexPrefetcher,
+        CpSharedKVIndexPrefetchHandle,
+        CpSharedKVMlaPrefetcher,
+        CpSharedKVMlaPrefetchHandle,
+    )
+    from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
+        build_in_seq_page_compute_owners,
+    )
+
+    page_size = 64
+    kv_dim = 656
+    index_page_nbytes = page_size * (128 + 4)
+    prefix_len, extend_len = 640, 200
+    prefix_pages = prefix_len // page_size
+    writers = [
+        int(o)
+        for o in build_in_seq_page_compute_owners(
+            extend_len=extend_len,
+            extend_prefix_len=prefix_len,
+            page_size=page_size,
+            cp_size=cp_size,
+        )
+    ]
+    pages = list(range(1, prefix_pages + 1))
+    next_page = prefix_pages + 1
+    for owner in writers:
+        while (next_page - 1) % cp_size != owner:
+            next_page += 1
+        pages.append(next_page)
+        next_page += 1
+    total_slots = len(pages)
+    logical_pages = torch.tensor([pages], dtype=torch.int64, device=device)
+    total_tokens = prefix_len + extend_len
+    locs = torch.tensor(
+        [
+            pages[i // page_size] * page_size + i % page_size
+            for i in range(total_tokens)
+        ],
+        dtype=torch.int64,
+        device=device,
+    )
+    physical_pages = (max(pages) - 1) // cp_size + 3
+    layout = CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=rank)
+
+    kv_cache = torch.zeros(
+        (physical_pages * page_size, 1, kv_dim), dtype=torch.uint8, device=device
+    )
+    page_buffer = torch.zeros(
+        (physical_pages, index_page_nbytes), dtype=torch.uint8, device=device
+    )
+    for logical_page in pages[:prefix_pages]:
+        if (logical_page - 1) % cp_size != rank:
+            continue
+        phys = (logical_page - 1) // cp_size + 1
+        kv_cache[phys * page_size : (phys + 1) * page_size] = (
+            (torch.arange(page_size * kv_dim, dtype=torch.int64) + logical_page * 53)
+            .remainder_(251)
+            .to(torch.uint8)
+            .view(page_size, 1, kv_dim)
+            .to(device)
+        )
+        page_buffer[phys] = (
+            (torch.arange(index_page_nbytes, dtype=torch.int64) + logical_page * 29)
+            .remainder_(251)
+            .to(torch.uint8)
+            .to(device)
+        )
+
+    cur_locs_all = [
+        pages[prefix_pages + i // page_size] * page_size + i % page_size
+        for i in range(extend_len)
+    ]
+    mine = [
+        i for i in range(extend_len) if writers[i // page_size] == rank
+    ]
+    current_locs = torch.tensor(
+        [cur_locs_all[i] for i in mine], dtype=torch.int64, device=device
+    )
+    rows = len(mine)
+    current_kv = (
+        (
+            torch.arange(rows, dtype=torch.int64).view(-1, 1, 1) * 11
+            + torch.arange(kv_dim, dtype=torch.int64).view(1, 1, -1)
+            + rank * 19
+        )
+        .remainder_(247)
+        .to(torch.uint8)
+        .to(device)
+    )
+    current_k = (
+        (
+            torch.arange(rows, dtype=torch.int64).view(-1, 1) * 23
+            + torch.arange(128, dtype=torch.int64).view(1, -1)
+            + rank * 7
+        )
+        .remainder_(245)
+        .to(torch.uint8)
+        .to(device)
+    )
+    current_scale = (
+        torch.arange(rows, dtype=torch.float32, device=device).view(-1, 1) * 0.5
+        + 2.0
+        + rank
+    )
+    zeros_req = torch.zeros(rows, dtype=torch.long, device=device)
+    loc_req_id = torch.zeros(total_tokens, dtype=torch.long, device=device)
+
+    kv_remap = rt.build_shared_token_kv_slot_remap(
+        kv_cache, locs.view(1, -1), logical_pages, layout, page_size
+    )
+    idx_remap = rt.build_shared_paged_buffer_slot_remap(
+        page_buffer, logical_pages, layout
+    )
+    spans = dict(
+        prefix_slot_spans=[(0, prefix_pages)],
+        current_slot_spans=[(prefix_pages, total_slots)],
+    )
+
+    # --- references via the legacy sync compose (pure NCCL, no barriers) ---
+    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
+        ref_kv, ref_locs = rt.materialize_prefix_and_reuse_current_kv_page_slots(
+            kv_cache=kv_cache,
+            logical_locs=locs,
+            current_kv_cache=current_kv,
+            current_locs=current_locs,
+            slot_remap=kv_remap,
+            layout=layout,
+            page_size=page_size,
+            prefix_pages=0,
+            loc_req_id=loc_req_id,
+            current_req_id=zeros_req,
+            layer_id=60,
+            **spans,
+        )
+        ref_idx, ref_pages = rt.materialize_prefix_and_reuse_current_index_page_slots(
+            page_buffer=page_buffer,
+            current_index_k=current_k,
+            current_index_scale=current_scale,
+            current_locs=current_locs,
+            slot_remap=idx_remap,
+            layout=layout,
+            page_size=page_size,
+            index_head_dim=128,
+            prefix_pages=0,
+            current_req_id=zeros_req,
+            layer_id=60,
+            **spans,
+        )
+    torch.cuda.synchronize()
+    dist.barrier()
+
+    # --- prefetchers with a hand-built prefix handle (prefix already
+    # materialized + reduced, exactly what start_next_layer_prefix yields) ---
+    mla = CpSharedKVMlaPrefetcher(
+        layout=layout,
+        page_size=page_size,
+        prefix_pages=prefix_pages,
+        slot_logical_pages=kv_remap.slot_logical_pages,
+        page_inverse=kv_remap.page_inverse,
+        dense_num_pages=kv_remap.dense_num_pages,
+        slot_remap=kv_remap,
+        symm_writers=writers,
+    )
+    dense_kv = kv_cache.new_zeros(
+        (kv_remap.dense_num_pages * page_size, 1, kv_dim)
+    )
+    rt.materialize_local_token_kv_page_slots_into(
+        kv_cache=kv_cache,
+        dense_kv_cache=dense_kv,
+        slot_logical_pages=kv_remap.slot_logical_pages,
+        layout=layout,
+        page_size=page_size,
+        start_slot=0,
+        end_slot=prefix_pages,
+    )
+    prefix_rows = rt.slot_range_to_token_slice(page_size, 0, prefix_pages)
+    dist.all_reduce(dense_kv[prefix_rows])
+    mla_event = torch.cuda.Event()
+    mla_event.record()
+    mla.handles[61] = CpSharedKVMlaPrefetchHandle(
+        layer_id=61, dense_kv_cache=dense_kv, prefix_rows=prefix_rows,
+        event=mla_event,
+    )
+
+    idx = CpSharedKVIndexPrefetcher(
+        layout=layout,
+        prefix_pages=prefix_pages,
+        slot_logical_pages=idx_remap.slot_logical_pages,
+        page_inverse=idx_remap.page_inverse,
+        dense_num_pages=idx_remap.dense_num_pages,
+        slot_remap=idx_remap,
+        symm_writers=writers,
+    )
+    dense_pb = page_buffer.new_zeros(
+        (idx_remap.dense_num_pages, index_page_nbytes)
+    )
+    rt.materialize_local_paged_buffer_page_slots_into(
+        page_buffer=page_buffer,
+        dense_page_buffer=dense_pb,
+        slot_logical_pages=idx_remap.slot_logical_pages,
+        layout=layout,
+        start_slot=0,
+        end_slot=prefix_pages,
+    )
+    dist.all_reduce(dense_pb[1 : prefix_pages + 1])
+    idx_event = torch.cuda.Event()
+    idx_event.record()
+    idx.handles[61] = CpSharedKVIndexPrefetchHandle(
+        layer_id=61, dense_page_buffer=dense_pb,
+        prefix_rows=slice(1, prefix_pages + 1), event=idx_event,
+    )
+    torch.cuda.synchronize()
+    dist.barrier()
+
+    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(
+        True
+    ), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True):
+        out_idx = idx.consume_prefix_with_current(
+            layer_id=61,
+            logical_pages=logical_pages,
+            current_index_k=current_k,
+            current_index_scale=current_scale,
+            current_locs=current_locs,
+            page_size=page_size,
+            index_head_dim=128,
+            current_req_id=zeros_req,
+        )
+        out_mla = mla.consume_prefix_with_current(
+            layer_id=61,
+            kv_cache=kv_cache,
+            logical_locs=locs,
+            current_kv_cache=current_kv,
+            current_locs=current_locs,
+            loc_req_id=loc_req_id,
+            current_req_id=zeros_req,
+        )
+    torch.cuda.synchronize()
+    assert out_idx is not None and out_mla is not None, (
+        f"rank{rank}: prefetch consume unexpectedly missed"
+    )
+    sym_idx, sym_pages = out_idx
+    sym_kv, sym_locs = out_mla
+    assert torch.equal(ref_pages, sym_pages), f"rank{rank}: prefetch index pages"
+    assert torch.equal(ref_locs, sym_locs), f"rank{rank}: prefetch mla locs"
+    if not torch.equal(ref_idx, sym_idx):
+        bad = (ref_idx != sym_idx).any(dim=-1).nonzero().reshape(-1)[:8]
+        raise AssertionError(
+            f"rank{rank}: prefetch index symm mismatch at pages {bad.cpu().tolist()}"
+        )
+    if not torch.equal(ref_kv, sym_kv):
+        diff = (ref_kv != sym_kv).any(dim=-1).any(dim=-1)
+        bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
+        raise AssertionError(
+            f"rank{rank}: prefetch mla symm mismatch at rows {bad} "
+            f"(of {int(diff.sum())})"
+        )
+
+
 def _compose(s, layer_id: int, *, writers: list[int] | None = None):
     return runtime.materialize_prefix_and_reuse_current_kv_page_slots(
         kv_cache=s["kv_cache"],
@@ -316,6 +636,65 @@ def main() -> None:
             "MiB, fill-to-staging + barrier + mega gather engaged; arena on/off)",
             flush=True,
         )
+
+    # ---- Phase 3: index sync compose, legacy vs symm (staging already
+    # registered by the token phase — the index compose never registers). ----
+    si = _build_index_scenario(rank, world, device, s)
+
+    def _compose_index(layer_id: int, *, writers=None):
+        return runtime.materialize_prefix_and_reuse_current_index_page_slots(
+            page_buffer=si["page_buffer"],
+            current_index_k=si["current_k"],
+            current_index_scale=si["current_scale"],
+            current_locs=si["current_locs"],
+            slot_remap=si["slot_remap"],
+            layout=s["layout"],
+            page_size=s["page_size"],
+            index_head_dim=128,
+            prefix_pages=0,
+            current_req_id=si["current_req_id"],
+            prefix_slot_spans=s["prefix_slot_spans"],
+            current_slot_spans=s["current_slot_spans"],
+            layer_id=layer_id,
+            current_page_writer_ranks=writers,
+        )
+
+    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
+        ref_idx, ref_dense_pages = _compose_index(40)
+    torch.cuda.synchronize()
+    dist.barrier()
+    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(
+        True
+    ), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True):
+        for layer_id in range(41, 45):
+            symm_idx, symm_dense_pages = _compose_index(
+                layer_id, writers=s["current_page_writer_ranks"]
+            )
+            torch.cuda.synchronize()
+            assert torch.equal(ref_dense_pages, symm_dense_pages), (
+                f"rank{rank} layer{layer_id}: index dense_pages mismatch"
+            )
+            if not torch.equal(ref_idx, symm_idx):
+                bad = (
+                    (ref_idx != symm_idx).any(dim=-1).nonzero().reshape(-1)[:8]
+                )
+                raise AssertionError(
+                    f"rank{rank} layer{layer_id}: index symm mismatch at pages "
+                    f"{bad.cpu().tolist()}"
+                )
+    dist.barrier()
+    if rank == 0:
+        print("PASS: index sync symm compose byte-identical (4 layers)", flush=True)
+
+    # ---- Phase 4+5: prefetch-path symm consume (MLA + index), bs=1. ----
+    _check_prefetch_symm(rank, world, device)
+    dist.barrier()
+    if rank == 0:
+        print(
+            "PASS: prefetch consume_prefix_with_current symm (MLA + index) "
+            "byte-identical to the sync compose",
+            flush=True,
+        )
     dist.destroy_process_group()
 
 

test/registered/unit/mem_cache/test_cp_shared_kv_compose.py

@@ -134,11 +134,11 @@ class TestComposePlanSymm(unittest.TestCase):
         self.assertEqual(plan.current_dense_pages.tolist(), [3, 4, 5])
         self.assertEqual(plan.symm_all_owner_ranks[:2].tolist(), [0, 1])
         self.assertEqual(plan.symm_all_owner_ranks[2:].tolist(), [11, 8, 11])
-        self.assertEqual(plan.symm_all_src_pages[2:].tolist(), [0, 1, 2])
+        self.assertEqual(plan.symm_all_src_pages[2:].tolist(), [1, 2, 3])
         # staging page inverse: current dense pages 3,4,5 -> staging slots
-        # 0,1,2; everything else (dummy, prefix pages) -> -1.
+        # 1,2,3 (slot 0 = dummy row); everything else -> -1.
         self.assertEqual(
-            plan.staging_page_inverse.tolist(), [[-1, -1, -1, 0, 1, 2]]
+            plan.staging_page_inverse.tolist(), [[-1, -1, -1, 1, 2, 3]]
         )
 
     def test_writer_count_mismatch_fails_fast(self):
@@ -310,7 +310,9 @@ class TestComposeStaging(unittest.TestCase):
         region_in_half = {}
         for kind in CpComposeStaging.KINDS:
             region_in_half[kind] = half
-            half += staging._align(capacity_pages * staging._page_nbytes[kind])
+            half += staging._align(
+                (capacity_pages + 1) * staging._page_nbytes[kind]
+            )
         for parity in (0, 1):
             for kind in CpComposeStaging.KINDS:
                 staging._region_offsets[(kind, parity)] = (
@@ -318,6 +320,8 @@ class TestComposeStaging(unittest.TestCase):
                 )
         staging._slab = torch.zeros(2 * half, dtype=torch.uint8)
         staging.peer_slab_bases = torch.full((8,), 1 << 20, dtype=torch.int64)
+        for key, offset in staging._region_offsets.items():
+            staging._region_peer_ptrs[key] = staging.peer_slab_bases + offset
         return staging, half
 
     def test_regions_are_disjoint_and_parity_halves_do_not_overlap(self):
@@ -326,12 +330,12 @@ class TestComposeStaging(unittest.TestCase):
         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))
+        # (index starts at the 256-aligned end of kv; +1 = dummy row 0).
+        self.assertEqual(idx0.data_ptr() - kv0.data_ptr(), staging._align(5 * 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)
+        self.assertEqual(kv0.numel(), 5 * 1000)
+        self.assertEqual(idx0.numel(), 5 * 300)
 
     def test_peer_region_ptrs_offset_matches_local_layout(self):
         staging, half = self._staging()
@@ -342,7 +346,7 @@ class TestComposeStaging(unittest.TestCase):
         self.assertTrue(
             torch.equal(
                 staging.peer_region_ptrs("index", 1),
-                base + half + staging._align(4 * 1000),
+                base + half + staging._align(5 * 1000),
             )
         )