Enable index partial-current reuse without replaying prefix materialize

The index path now mirrors the target MLA partial-current contract: prefetched or synchronously materialized prefix pages are composed with valid current index K/scale rows in slot-dense page buffers. Current-only batches keep the compact current-index fast path, while partial cache-hit batches share one composed dense index buffer across the in-seq prev/next topk pair.\n\nThe prefetch consume path remaps through the slot page inverse instead of treating the slot-dense buffer as physical-pool capacity, and current-index quantization uses valid extend rows so padded out_cache_loc does not disable reuse.\n\nConstraint: CP shared KV remains page-slot based; padding rows must stay invisible to attention/index semantics\nConstraint: Draft/EAGLE partial-current reuse remains guarded by should_reuse_current_extend_kv\nRejected: Replace prefix all-reduce with all-gather | NCCL all-gather still uses SM and would require an additional compose/scatter step\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce current-only gating for index reuse; partial target cache hits must compose prefix + valid current rows\nTested: Local py_compile for touched Python files\nTested: g0034 sglang-glm5-dev-2 PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 77 passed, 5 warnings, 2 subtests passed\nNot-tested: Full ETE traffic with latest commit; CUDA perf impact of index partial-current prefetch under production load

docs/advanced_features/nsa_prefill_cp_phase6_index_materialize_reuse_plan.md

@@ -91,3 +91,36 @@ PYTHONPATH=python python3 -m pytest \
   test/registered/unit/layers/test_nsa_cp_utils.py \
   test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -q
 ```
+
+## 2026-05-31 update: index partial/current reuse
+
+The index path now follows the same target-model partial-current contract as MLA:
+
+```text
+page-aligned cached prefix index pages
++ valid current index K/scale rows copied into their slot-dense suffix pages
++ page-tail slack left zero and invisible through valid seq/page-table lengths
+```
+
+Implementation notes:
+
+- `Indexer._can_reuse_current_index_kv()` uses `should_reuse_current_extend_kv()` instead of the old current-only gate, so target cache-hit suffixes can reuse freshly computed current index K/scale.
+- Current-only batches still use the compact `current_index_kv` path and avoid page-table materialization.
+- Partial cache-hit batches materialize the prefix once for the prev/next in-seq split pair, then pass the composed dense index buffer to both topk calls.
+- If the async index prefetcher has a ready prefix handle, `CpSharedKVIndexPrefetcher.consume_prefix_with_current()` consumes it and only fills current rows into the suffix page slots.
+- If no prefetch handle exists, `materialize_prefix_and_reuse_current_index_page_slots()` synchronously materializes only prefix slots, all-reduces that prefix range, and fills current rows locally.
+- `forward_cuda()` quantizes only `valid_current_rows = extend_seq_lens_cpu[0]` current rows. This preserves current reuse when `out_cache_loc` is physically padded for page alignment.
+
+Debug contract correction:
+
+- The prefetch consume path must remap `logical_pages` through the slot `page_inverse` directly. It must not validate these logical pages against `dense_page_buffer.shape[0]` as if the dense buffer were the physical pool capacity; that debug-only check can reject valid high logical page ids that are intentionally packed into a small slot-dense buffer.
+
+Verification:
+
+```bash
+# g0034 / sglang-glm5-dev-2
+cd /sgl-workspace/sglang-tai
+PYTHONPATH=python python -m pytest -q \
+  test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+# 77 passed, 5 warnings, 2 subtests passed
+```

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

@@ -19,6 +19,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     cp_shared_kv_mla_prefetch_should_log_layer,
     filter_locs_mappable_to_physical_pool,
     filter_pages_mappable_to_physical_pool,
+    fill_current_index_page_slots,
     fill_current_kv_page_slots_and_remap_locs,
     get_or_build_shared_paged_buffer_slot_remap,
     get_or_build_shared_token_kv_slot_remap,
@@ -1469,6 +1470,104 @@ class CpSharedKVIndexPrefetcher:
         )
         return dense_page_buffer, dense_pages
 
+    def consume_prefix_with_current(
+        self,
+        *,
+        layer_id: int,
+        logical_pages: torch.Tensor,
+        current_index_k: torch.Tensor,
+        current_index_scale: torch.Tensor,
+        current_locs: torch.Tensor,
+        page_size: int,
+        index_head_dim: int,
+    ) -> Optional[tuple[torch.Tensor, torch.Tensor]]:
+        if self.disabled:
+            self._log_layer(
+                layer_id,
+                "index_consume_prefix_current_skip reason=disabled layer=%s",
+                layer_id,
+            )
+            return None
+
+        self._log_layer(
+            layer_id,
+            "index_consume_prefix_current_enter layer=%s prefix_pages=%s "
+            "total_slots=%s handles=%s",
+            layer_id,
+            self.prefix_pages,
+            self.total_slots,
+            _debug_handle_keys(layer_id, self.handles),
+        )
+        handle = self.handles.get(layer_id)
+        if handle is None:
+            self._log_layer(
+                layer_id,
+                "index_consume_prefix_current_miss layer=%s",
+                layer_id,
+            )
+            return None
+        if handle.event is None:
+            self.launch_pending_reduce()
+            handle = self.handles.get(layer_id)
+            if handle is None or handle.event is None:
+                self._log_layer(
+                    layer_id,
+                    "index_consume_prefix_current_miss "
+                    "reason=prefix_reduce_not_ready layer=%s",
+                    layer_id,
+                )
+                return None
+        handle = self.handles.pop(layer_id)
+        if self.pending_attention_handle is handle:
+            self.pending_attention_handle = None
+        if handle.layer_id != layer_id:
+            self.disabled = True
+            self._log_layer(
+                layer_id,
+                "index_consume_prefix_current_skip reason=layer_mismatch "
+                "expected=%s actual=%s",
+                layer_id,
+                handle.layer_id,
+            )
+            return None
+
+        consume_cpu = _cpu_timing_start()
+        wait_cpu = _cpu_timing_start()
+        torch.cuda.current_stream().wait_event(handle.event)
+        wait_ms = _cpu_timing_ms(wait_cpu)
+        dense_page_buffer = handle.dense_page_buffer
+
+        remap_cpu = _cpu_timing_start()
+        dense_pages = remap_logical_pages_to_slot_dense_pages(
+            logical_pages,
+            page_inverse=self.page_inverse,
+        )
+        dense_page_buffer = fill_current_index_page_slots(
+            dense_page_buffer=dense_page_buffer,
+            current_index_k=current_index_k,
+            current_index_scale=current_index_scale,
+            current_locs=current_locs,
+            page_inverse=self.page_inverse,
+            page_size=page_size,
+            index_head_dim=index_head_dim,
+        )
+        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 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, dense_pages
+
     def start_next_layer_prefix(
         self,
         *,

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

@@ -813,6 +813,103 @@ def fill_current_kv_page_slots_and_remap_locs(
     return dense_kv_cache, mixed_locs, current_mask
 
 
+def fill_current_index_page_slots(
+    *,
+    dense_page_buffer: torch.Tensor,
+    current_index_k: torch.Tensor,
+    current_index_scale: torch.Tensor,
+    current_locs: torch.Tensor,
+    page_inverse: torch.Tensor,
+    page_size: int,
+    index_head_dim: int,
+) -> torch.Tensor:
+    """Fill current index K/scale rows into a slot-dense page buffer.
+
+    The dense page buffer already has one row per logical page slot and page 0 is
+    the dummy page.  Current rows are copied into the row/offset selected by their
+    logical token loc; tail-page slack remains zero and is invisible because the
+    index kernels still receive the valid sequence length separately.
+    """
+
+    current_locs = current_locs.reshape(-1)
+    current_rows = int(current_locs.numel())
+    if current_rows == 0:
+        return dense_page_buffer
+    if (
+        int(current_index_k.shape[0]) < current_rows
+        or int(current_index_scale.shape[0]) < current_rows
+    ):
+        raise ValueError(
+            "CP shared KV index current-slot fill got fewer current rows than "
+            f"locs: k_rows={int(current_index_k.shape[0])} "
+            f"scale_rows={int(current_index_scale.shape[0])} locs={current_rows}"
+        )
+
+    current_index_k = current_index_k[:current_rows].contiguous()
+    current_index_scale = current_index_scale[:current_rows].contiguous()
+    k_bytes = current_index_k.reshape(current_rows, -1).view(torch.uint8)
+    if int(k_bytes.shape[1]) != int(index_head_dim):
+        raise ValueError(
+            "CP shared KV index current-slot fill got unexpected K width: "
+            f"k_bytes_per_token={int(k_bytes.shape[1])} index_head_dim={index_head_dim}"
+        )
+    scale_bytes = current_index_scale.reshape(current_rows, -1).view(torch.uint8)
+    scale_bytes = scale_bytes.reshape(current_rows, -1)
+
+    page_stride = int(dense_page_buffer.shape[1])
+    scale_offset = int(page_size) * int(index_head_dim)
+    if scale_offset + int(page_size) * int(scale_bytes.shape[1]) > page_stride:
+        raise ValueError(
+            "CP shared KV index current-slot fill got incompatible page buffer: "
+            f"page_stride={page_stride} scale_offset={scale_offset} "
+            f"scale_bytes_per_token={int(scale_bytes.shape[1])} page_size={page_size}"
+        )
+
+    current_pages = torch.div(current_locs, page_size, rounding_mode="floor")
+    valid_pages = (current_locs >= 0) & (current_pages >= 0) & (
+        current_pages < int(page_inverse.numel())
+    )
+    safe_pages = torch.clamp(
+        current_pages,
+        min=0,
+        max=max(int(page_inverse.numel()) - 1, 0),
+    )
+    dense_pages = page_inverse[safe_pages.to(torch.long)].to(torch.long)
+    valid_rows = valid_pages & (dense_pages > 0)
+    if not torch.any(valid_rows):
+        return dense_page_buffer
+
+    valid_dense_pages = dense_pages[valid_rows]
+    valid_offsets = torch.remainder(current_locs[valid_rows], page_size).to(torch.long)
+    flat_buffer = dense_page_buffer.reshape(-1)
+
+    k_cols = torch.arange(
+        int(index_head_dim),
+        dtype=torch.long,
+        device=dense_page_buffer.device,
+    )
+    k_indices = (
+        valid_dense_pages[:, None] * page_stride
+        + valid_offsets[:, None] * int(index_head_dim)
+        + k_cols[None, :]
+    )
+    flat_buffer[k_indices.reshape(-1)] = k_bytes[valid_rows].reshape(-1)
+
+    s_cols = torch.arange(
+        int(scale_bytes.shape[1]),
+        dtype=torch.long,
+        device=dense_page_buffer.device,
+    )
+    s_indices = (
+        valid_dense_pages[:, None] * page_stride
+        + scale_offset
+        + valid_offsets[:, None] * int(scale_bytes.shape[1])
+        + s_cols[None, :]
+    )
+    flat_buffer[s_indices.reshape(-1)] = scale_bytes[valid_rows].reshape(-1)
+    return dense_page_buffer
+
+
 def _copy_tai_dense_slot_range_body(
     *,
     tai_dense_kv_cache: torch.Tensor,
@@ -2058,6 +2155,62 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
     return mixed_kv_cache, mixed_locs
 
 
+def materialize_prefix_and_reuse_current_index_page_slots(
+    *,
+    page_buffer: torch.Tensor,
+    current_index_k: torch.Tensor,
+    current_index_scale: torch.Tensor,
+    current_locs: torch.Tensor,
+    slot_remap: SharedPagedBufferSlotRemap,
+    layout: CpSharedKVLayout,
+    page_size: int,
+    index_head_dim: int,
+    prefix_pages: int,
+    layer_id: int | None = None,
+    nvtx_source: str = "index.partial_current_sync",
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Synchronously compose prefix index materialization with current index rows."""
+
+    total_slots = int(slot_remap.slot_logical_pages.numel())
+    if prefix_pages < 0 or prefix_pages > total_slots:
+        raise ValueError(
+            "Invalid CP shared KV index partial-current prefix range: "
+            f"prefix_pages={prefix_pages} total_slots={total_slots}"
+        )
+
+    dense_page_buffer = page_buffer.new_zeros(
+        (slot_remap.dense_num_pages, *page_buffer.shape[1:])
+    )
+    materialize_local_paged_buffer_page_slots_into(
+        page_buffer=page_buffer,
+        dense_page_buffer=dense_page_buffer,
+        slot_logical_pages=slot_remap.slot_logical_pages,
+        layout=layout,
+        start_slot=0,
+        end_slot=prefix_pages,
+    )
+    prefix_rows = slot_range_to_page_slice(0, prefix_pages)
+    _all_reduce_materialized_buffer_range(
+        dense_page_buffer,
+        layout.cp_size,
+        prefix_rows.start,
+        prefix_rows.stop,
+        nvtx_source=nvtx_source,
+        nvtx_layer_id=layer_id,
+        nvtx_cp_rank=layout.cp_rank,
+    )
+    dense_page_buffer = fill_current_index_page_slots(
+        dense_page_buffer=dense_page_buffer,
+        current_index_k=current_index_k,
+        current_index_scale=current_index_scale,
+        current_locs=current_locs,
+        page_inverse=slot_remap.page_inverse,
+        page_size=page_size,
+        index_head_dim=index_head_dim,
+    )
+    return dense_page_buffer, slot_remap.dense_pages
+
+
 def slot_range_to_token_slice(
     page_size: int,
     start_slot: int,

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

@@ -17,7 +17,6 @@ from sglang.srt.layers.attention.nsa import index_buf_accessor
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     cp_shared_kv_debug_enabled,
     cp_shared_kv_debug_log,
-    cp_shared_kv_current_reuse_enabled,
     cp_shared_kv_mla_prefetch_log,
     cp_shared_kv_mla_prefetch_log_enabled,
     cp_shared_kv_mla_prefetch_should_log_layer,
@@ -26,7 +25,9 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     get_or_build_shared_paged_buffer_slot_remap,
     is_current_only_extend_batch,
     log_cp_draft_shared_kv_debug,
+    materialize_prefix_and_reuse_current_index_page_slots,
     materialize_shared_paged_buffer,
+    should_reuse_current_extend_kv,
     tensor_debug_checksum,
     tensor_debug_summary,
     try_tai_prepare_cp_mqa_index,
@@ -302,6 +303,7 @@ class Indexer(MultiPlatformOp):
         forward_batch: ForwardBatch,
         layer_id: int,
         logical_page_table: torch.Tensor,
+        current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         index_buffer = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer(
             layer_id=layer_id
@@ -314,6 +316,100 @@ class Indexer(MultiPlatformOp):
         index_prefetcher = getattr(
             forward_batch, "cp_shared_kv_index_prefetcher", None
         )
+        if current_index_kv is not None:
+            page_size = int(forward_batch.token_to_kv_pool.page_size)
+            prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None)
+            extend_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
+            prefix_lens = (
+                [int(x) for x in prefix_lens_cpu]
+                if prefix_lens_cpu is not None
+                else None
+            )
+            extend_lens = (
+                [int(x) for x in extend_lens_cpu]
+                if extend_lens_cpu is not None
+                else None
+            )
+            if (
+                prefix_lens_cpu is None
+                or len(prefix_lens_cpu) != 1
+                or int(prefix_lens_cpu[0]) <= 0
+                or int(prefix_lens_cpu[0]) % page_size != 0
+            ):
+                raise RuntimeError(
+                    "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
+                    "CP shared KV index partial-current compose requires one "
+                    "positive page-aligned prefix. "
+                    f"cp_rank={layout.cp_rank} layer_id={layer_id} "
+                    f"prefix_lens={prefix_lens} extend_lens={extend_lens} "
+                    f"logical_page_table_shape={tuple(logical_page_table.shape)} "
+                    f"page_size={page_size}"
+                )
+            current_locs = forward_batch.out_cache_loc
+            if extend_lens_cpu is not None and len(extend_lens_cpu) == 1:
+                valid_current_rows = int(extend_lens_cpu[0])
+                if (
+                    valid_current_rows > 0
+                    and valid_current_rows < int(current_locs.numel())
+                    and valid_current_rows <= int(current_index_kv[0].shape[0])
+                    and valid_current_rows <= int(current_index_kv[1].shape[0])
+                ):
+                    current_locs = current_locs[:valid_current_rows]
+                    current_index_kv = (
+                        current_index_kv[0][:valid_current_rows],
+                        current_index_kv[1][:valid_current_rows],
+                    )
+            prefix_pages = int(prefix_lens_cpu[0]) // page_size
+            if index_prefetcher is not None:
+                prefetched = index_prefetcher.consume_prefix_with_current(
+                    layer_id=layer_id,
+                    logical_pages=logical_page_table,
+                    current_index_k=current_index_kv[0],
+                    current_index_scale=current_index_kv[1],
+                    current_locs=current_locs,
+                    page_size=page_size,
+                    index_head_dim=forward_batch.token_to_kv_pool.index_head_dim,
+                )
+                if prefetched is not None:
+                    return prefetched
+            slot_remap = get_or_build_shared_paged_buffer_slot_remap(
+                forward_batch,
+                page_buffer=index_buffer,
+                logical_pages=logical_page_table,
+                layout=layout,
+            )
+            materialized, dense_pages = (
+                materialize_prefix_and_reuse_current_index_page_slots(
+                    page_buffer=index_buffer,
+                    current_index_k=current_index_kv[0],
+                    current_index_scale=current_index_kv[1],
+                    current_locs=current_locs,
+                    slot_remap=slot_remap,
+                    layout=layout,
+                    page_size=page_size,
+                    index_head_dim=forward_batch.token_to_kv_pool.index_head_dim,
+                    prefix_pages=prefix_pages,
+                    layer_id=layer_id,
+                )
+            )
+            if (
+                cp_shared_kv_mla_prefetch_log_enabled()
+                and cp_shared_kv_mla_prefetch_should_log_layer(layer_id)
+            ):
+                cp_shared_kv_mla_prefetch_log(
+                    "index_partial_current_sync_compose cp_rank=%s layer=%s "
+                    "prefix_lens=%s extend_lens=%s prefix_pages=%s "
+                    "current_rows=%s dense_pages=%s",
+                    layout.cp_rank,
+                    layer_id,
+                    prefix_lens,
+                    extend_lens,
+                    prefix_pages,
+                    int(current_index_kv[0].shape[0]),
+                    int(materialized.shape[0]),
+                )
+            return materialized, dense_pages
+
         if index_prefetcher is not None:
             prefetched = index_prefetcher.consume(
                 layer_id=layer_id,
@@ -461,12 +557,11 @@ class Indexer(MultiPlatformOp):
 
     def _can_reuse_current_index_kv(self, forward_batch: ForwardBatch) -> bool:
         return (
-            cp_shared_kv_current_reuse_enabled()
-            and forward_batch.uses_cp_shared_kv
+            forward_batch.uses_cp_shared_kv
             and self.nsa_enable_prefill_cp
             and forward_batch.nsa_cp_metadata is not None
             and is_nsa_prefill_cp_in_seq_split()
-            and is_current_only_extend_batch(forward_batch)
+            and should_reuse_current_extend_kv(forward_batch)
             and forward_batch.hisparse_coordinator is None
             and _is_cuda
             and not _is_fp8_fnuz
@@ -1283,7 +1378,21 @@ class Indexer(MultiPlatformOp):
 
         shared_index_buffer = None
         shared_block_tables = None
-        if current_index_kv is None:
+        current_index_kv_for_topk = current_index_kv
+        if current_index_kv is not None and not is_current_only_extend_batch(
+            forward_batch
+        ):
+            current_index_kv_for_topk = None
+            shared_block_tables = metadata.get_page_table_64()
+            shared_index_buffer, shared_block_tables = (
+                self._maybe_materialize_shared_index_buffer(
+                    forward_batch,
+                    layer_id,
+                    shared_block_tables,
+                    current_index_kv=current_index_kv,
+                )
+            )
+        elif current_index_kv is None:
             shared_block_tables = metadata.get_page_table_64()
             shared_index_buffer, shared_block_tables = (
                 self._maybe_materialize_shared_index_buffer(
@@ -1301,7 +1410,7 @@ class Indexer(MultiPlatformOp):
             metadata,
             kv_len_prev,
             actual_seq_q_prev,
-            current_index_kv=current_index_kv,
+            current_index_kv=current_index_kv_for_topk,
             shared_index_buffer=shared_index_buffer,
             shared_block_tables=shared_block_tables,
             actual_seq_q_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_prev_tensor,
@@ -1316,7 +1425,7 @@ class Indexer(MultiPlatformOp):
             metadata,
             kv_len_next,
             actual_seq_q_next,
-            current_index_kv=current_index_kv,
+            current_index_kv=current_index_kv_for_topk,
             shared_index_buffer=shared_index_buffer,
             shared_block_tables=shared_block_tables,
             actual_seq_q_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_next_tensor,
@@ -1674,9 +1783,18 @@ class Indexer(MultiPlatformOp):
 
         current_index_kv = None
         if self._can_reuse_current_index_kv(forward_batch):
-            if key.shape[0] == forward_batch.out_cache_loc.numel():
+            extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
+            valid_current_rows = int(forward_batch.out_cache_loc.numel())
+            if extend_seq_lens_cpu is not None and len(extend_seq_lens_cpu) == 1:
+                valid_current_rows = min(
+                    int(extend_seq_lens_cpu[0]),
+                    valid_current_rows,
+                )
+            if key.shape[0] >= valid_current_rows:
                 current_k_fp8, current_k_scale = act_quant(
-                    key.contiguous(), self.block_size, self.scale_fmt
+                    key[:valid_current_rows].contiguous(),
+                    self.block_size,
+                    self.scale_fmt,
                 )
                 current_index_kv = (
                     current_k_fp8.contiguous(),
@@ -1690,14 +1808,25 @@ class Indexer(MultiPlatformOp):
                         if forward_batch.cp_shared_kv_layout is not None
                         else None,
                         layer_id,
-                        tensor_debug_summary(forward_batch.out_cache_loc),
+                        tensor_debug_summary(
+                            forward_batch.out_cache_loc[:valid_current_rows]
+                        ),
                         tensor_debug_checksum(current_index_kv[0]),
                         tensor_debug_checksum(current_index_kv[1]),
                     )
-            elif cp_shared_kv_debug_enabled():
-                raise RuntimeError(
-                    "CP shared KV current index reuse shape mismatch: "
-                    f"key_tokens={key.shape[0]} out_cache_loc={forward_batch.out_cache_loc.numel()}"
+            else:
+                _log_cp_shared_kv_index_prefetch_fallback(
+                    "current_reuse_shape_mismatch",
+                    "NSA index current reuse skipped because key has fewer rows "
+                    "than valid current locs. cp_rank=%s layer=%s key_tokens=%s "
+                    "valid_current_rows=%s out_cache_loc=%s",
+                    forward_batch.cp_shared_kv_layout.cp_rank
+                    if forward_batch.cp_shared_kv_layout is not None
+                    else None,
+                    layer_id,
+                    int(key.shape[0]),
+                    valid_current_rows,
+                    int(forward_batch.out_cache_loc.numel()),
                 )
 
         if _is_cuda or _is_hip:

test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py

@@ -1021,6 +1021,182 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         self.assertLess(valid_tokens, padded_pages * page_size)
         self.assertEqual(valid_tokens, 10)
 
+    def test_index_partial_current_sync_compose_fills_current_page_slots(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+
+        page_size = 4
+        index_head_dim = 4
+        scale_bytes = 4
+        page_bytes = page_size * index_head_dim + page_size * scale_bytes
+        layout = CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0)
+        page_buffer = torch.zeros((6, page_bytes), dtype=torch.uint8)
+        page_buffer[1] = torch.arange(10, 10 + page_bytes, dtype=torch.uint8)
+        page_buffer[2] = 99
+        logical_pages = torch.tensor([[1, 2]], dtype=torch.int64)
+        slot_remap = runtime.build_shared_paged_buffer_slot_remap(
+            page_buffer,
+            logical_pages,
+            layout,
+        )
+        current_k = torch.tensor(
+            [[1, 2, 3, 4], [5, 6, 7, 8]],
+            dtype=torch.uint8,
+        )
+        current_scale = torch.tensor([[1.25], [2.5]], dtype=torch.float32)
+
+        with patch.object(
+            runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce
+        ):
+            dense_page_buffer, dense_pages = (
+                runtime.materialize_prefix_and_reuse_current_index_page_slots(
+                    page_buffer=page_buffer,
+                    current_index_k=current_k,
+                    current_index_scale=current_scale,
+                    current_locs=torch.tensor([8, 9], dtype=torch.int64),
+                    slot_remap=slot_remap,
+                    layout=layout,
+                    page_size=page_size,
+                    index_head_dim=index_head_dim,
+                    prefix_pages=1,
+                    layer_id=2,
+                )
+            )
+
+        scale_offset = page_size * index_head_dim
+        self.assertEqual(dense_pages.tolist(), [[1, 2]])
+        self.assertTrue(torch.equal(dense_page_buffer[1], page_buffer[1]))
+        self.assertTrue(torch.equal(dense_page_buffer[2, 0:4], current_k[0]))
+        self.assertTrue(torch.equal(dense_page_buffer[2, 4:8], current_k[1]))
+        self.assertTrue(
+            torch.equal(
+                dense_page_buffer[2, 8:scale_offset],
+                torch.zeros(8, dtype=torch.uint8),
+            )
+        )
+        self.assertTrue(
+            torch.equal(
+                dense_page_buffer[2, scale_offset : scale_offset + 4],
+                current_scale[0].view(torch.uint8),
+            )
+        )
+        self.assertTrue(
+            torch.equal(
+                dense_page_buffer[2, scale_offset + 4 : scale_offset + 8],
+                current_scale[1].view(torch.uint8),
+            )
+        )
+        self.assertTrue(
+            torch.equal(
+                dense_page_buffer[2, scale_offset + 8 :],
+                torch.zeros(8, dtype=torch.uint8),
+            )
+        )
+
+    def test_index_prefetch_partial_current_compose_fills_current_page_slots(self):
+        from sglang.srt.environ import envs
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
+        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+
+        class FakeCurrentStream:
+            def __init__(self):
+                self.events = []
+
+            def wait_event(self, event):
+                self.events.append(event)
+
+        page_size = 4
+        index_head_dim = 4
+        scale_bytes = 4
+        page_bytes = page_size * index_head_dim + page_size * scale_bytes
+        dense_page_buffer = torch.zeros((3, page_bytes), dtype=torch.uint8)
+        dense_page_buffer[1] = torch.arange(10, 10 + page_bytes, dtype=torch.uint8)
+        fake_event = object()
+        prefetcher = prefetch.CpSharedKVIndexPrefetcher(
+            layout=CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0),
+            prefix_pages=1,
+            slot_logical_pages=torch.tensor([1, 20], dtype=torch.int64),
+            page_inverse=torch.tensor(
+                [-1, 1] + [-1] * 18 + [2],
+                dtype=torch.int64,
+            ),
+            dense_num_pages=3,
+            stream=object(),
+        )
+        handle = prefetch.CpSharedKVIndexPrefetchHandle(
+            layer_id=1,
+            dense_page_buffer=dense_page_buffer,
+            prefix_rows=slice(1, 2),
+            event=fake_event,
+        )
+        prefetcher.handles[1] = handle
+        prefetcher.pending_attention_handle = handle
+        current_k = torch.tensor(
+            [[11, 12, 13, 14], [15, 16, 17, 18]],
+            dtype=torch.uint8,
+        )
+        current_scale = torch.tensor([[3.25], [4.5]], dtype=torch.float32)
+
+        with envs.SGLANG_DEBUG_CP_SHARED_KV.override(True), patch.object(
+            prefetch.torch.cuda,
+            "current_stream",
+            return_value=FakeCurrentStream(),
+        ):
+            mixed_buffer, dense_pages = prefetcher.consume_prefix_with_current(
+                layer_id=1,
+                logical_pages=torch.tensor([[1, 20]], dtype=torch.int64),
+                current_index_k=current_k,
+                current_index_scale=current_scale,
+                current_locs=torch.tensor([80, 81], dtype=torch.int64),
+                page_size=page_size,
+                index_head_dim=index_head_dim,
+            )
+
+        scale_offset = page_size * index_head_dim
+        self.assertEqual(dense_pages.tolist(), [[1, 2]])
+        self.assertTrue(torch.equal(mixed_buffer[2, 0:4], current_k[0]))
+        self.assertTrue(torch.equal(mixed_buffer[2, 4:8], current_k[1]))
+        self.assertTrue(
+            torch.equal(
+                mixed_buffer[2, scale_offset : scale_offset + 4],
+                current_scale[0].view(torch.uint8),
+            )
+        )
+        self.assertEqual(prefetcher.handles, {})
+        self.assertIsNone(prefetcher.pending_attention_handle)
+
+    def test_index_current_reuse_gate_uses_partial_current_contract(self):
+        from pathlib import Path
+
+        source = (
+            Path(__file__).resolve().parents[4]
+            / "python/sglang/srt/layers/attention/nsa/nsa_indexer.py"
+        ).read_text()
+        start = source.index("    def _can_reuse_current_index_kv")
+        end = source.index("    @contextlib.contextmanager", start)
+        method_source = source[start:end]
+
+        self.assertIn("should_reuse_current_extend_kv(forward_batch)", method_source)
+        self.assertNotIn("is_current_only_extend_batch(forward_batch)", method_source)
+
+    def test_index_current_reuse_prepare_accepts_padded_out_cache_loc(self):
+        from pathlib import Path
+
+        source = (
+            Path(__file__).resolve().parents[4]
+            / "python/sglang/srt/layers/attention/nsa/nsa_indexer.py"
+        ).read_text()
+        start = source.index("        if self._can_reuse_current_index_kv")
+        end = source.index("        if _is_cuda or _is_hip:", start)
+        prepare_source = source[start:end]
+
+        self.assertIn("valid_current_rows", prepare_source)
+        self.assertIn("key[:valid_current_rows]", prepare_source)
+        self.assertNotIn(
+            "key.shape[0] == forward_batch.out_cache_loc.numel()",
+            prepare_source,
+        )
+
     def test_materialize_local_token_kv_pages(self):
         from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
             build_dense_page_remap,