Reduce CP shared KV overhead without changing ownership semantics

The shared-KV path now keeps more CP metadata on-device and reuses
physical out-cache locations across MLA and NSA index writes, so each
layer avoids repeating logical-to-physical remaps. The in-seq CP
all-gather rerange path now delegates to tai-kernel when available and
falls back to the existing torch split/cat path with an explicit log.

This also extends the Phase8 prefetch machinery to cover shared KV
materialization metadata and keeps debug/fallback behavior gated so the
fast path is not polluted by diagnostic checks.

Constraint: Custom CP kernels must live in tai-kernel and be imported lazily from SGLang
Constraint: Decode does not use CP; these changes target NSA prefill CP in-seq-split shared KV
Rejected: Recompute physical local cache locations separately for MLA and index writes | repeats the same remap work every layer
Rejected: Keep the in-seq rerange Triton code inline in SGLang | duplicates kernel ownership and blocks tai-kernel reuse
Confidence: medium
Scope-risk: moderate
Directive: Keep CP collective ordering identical across ranks; do not add rank-local fallback decisions inside shared KV materialize paths
Tested: Remote g0034 container py_compile for modified SGLang/tai-kernel files; remote pytest test/registered/unit/layers/test_nsa_cp_utils.py passed with 24 tests
Not-tested: Full multi-node GLM5 prefill/decode throughput after the final commit boundary

docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md

@@ -1,8 +1,8 @@
-# NSA Prefill CP Phase 8: MLA prefix prefetch
+# NSA Prefill CP Phase 8: MLA / index prefix prefetch
 
-Phase 8 的目标是在不改变 Phase 2-7 shared KV 语义的前提下，为 chunked prefill / radix cache hit 场景引入 **MLA KV one-layer-ahead prefix prefetch**，把历史 prefix 的 shared KV materialize 从当前层 attention 前的同步阻塞路径里移出，并把等待点延迟到下一层真正消费 prefetched KV 时。
+Phase 8 的目标是在不改变 Phase 2-7 shared KV 语义的前提下，为 chunked prefill / radix cache hit 场景引入 **MLA KV 与 NSA index K/scale 的 one-layer-ahead prefix prefetch**，把历史 prefix 的 shared KV/index materialize 从当前层同步阻塞路径里移出，并把等待点延迟到下一层真正消费 prefetched buffer 时。
 
-本阶段只做 MLA KV prefix prefetch。暂时不做 index K/scale prefetch，不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND` 这类选择型环境变量，也不保留 `SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION` 这类 wait 策略开关。
+本阶段不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND` 这类选择型环境变量，也不保留 `SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION` 这类 wait 策略开关。`SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1` 现在同时打开 MLA KV prefix prefetch 和 NSA index K/scale prefix prefetch。
 
 ---
 
@@ -12,15 +12,16 @@ Phase 8 的目标是在不改变 Phase 2-7 shared KV 语义的前提下，为 ch
 Phase 8 已按 deferred-consume 策略实现：
 
 ```text
+Layer L index/topk 后启动 Layer L+1 prefix index K/scale prefetch
 Layer L attention 前启动 Layer L+1 prefix MLA KV prefetch
-Layer L forward_extend 返回时不等待该 prefetch
-Layer L+1 consume prefetched KV 时 wait event，再补齐 suffix/current pages
+Layer L attention 返回时不强制等待 prefetch
+Layer L+1 consume prefetched buffer 时 wait event，再补齐 suffix/current pages
 ```
 
 保留的环境变量：
 
 ```text
-SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1   # 生产开关，默认关闭
+SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1   # 生产开关，默认关闭；同时控制 MLA KV 与 index K/scale prefetch
 SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=0/1      # 调试日志，只打印 probe layer，默认关闭
 ```
 
@@ -48,13 +49,13 @@ owner-sharded physical MLA KV on each CP rank
   -> existing NSA attention kernel
 ```
 
-在 chunked prefill 的第二个 chunk 及之后，或者 radix cache 命中时，`extend_prefix_len > 0`。这部分 prefix KV 已经在之前的 chunk/request 中写入 persistent KV pool；当前层 attention 仍会同步 materialize 整个可见 KV，包括历史 prefix 和当前 suffix。
+在 chunked prefill 的第二个 chunk 及之后，或者 radix cache 命中时，`extend_prefix_len > 0`。这部分 prefix KV 与 index K/scale 已经在之前的 chunk/request 中写入 persistent pool；当前层 index/topk 和 attention 仍会同步 materialize 整个可见范围，包括历史 prefix 和当前 suffix。
 
 Phase 8 的机会是：
 
 ```text
 历史 prefix 部分已经存在，可以提前为下一层 materialize。
-当前 suffix 部分必须等下一层 prepare 写入后才能 materialize。
+当前 suffix 部分必须等下一层 prepare/index write 写入后才能 materialize。
 ```
 
 因此 Phase 8 不做跨 chunk 的 dense KV cache 复用，而是做 **每次 forward 内、相邻 layer 之间的一层提前预取**。
@@ -173,12 +174,12 @@ _get_topk_in_seq_cp_pair(...)
   -> materialize_shared_paged_buffer(...)
 ```
 
-Phase 8 暂时不碰 index path。原因：
+Index K/scale prefetch 在 MLA KV prefetch 之后补齐。当前约束：
 
 1. index materialize 已经完成 Phase 6 的一次合并；
 2. indexer 在 MLA attention 之前执行，调度窗口不同；
-3. 本阶段目标是先验证 MLA prefix prefetch 是否能显著隐藏最大块的 KV materialize；
-4. 避免同时改变 topk/index 和 attention KV 两条路径，降低正确性风险。
+3. 当前实现先落地 MLA KV，再补 index K/scale；两者共用同一个 page-aligned slot layout；
+4. index K/scale prefetch 只覆盖 PAGED topk 路径，RAGGED 仍不进入本阶段。
 
 ---
 
@@ -192,9 +193,16 @@ Phase 8 暂时不碰 index path。原因：
 Layer L attention 计算期间，
 提前为 Layer L+1 materialize prefix MLA KV。
 
+Layer L topk/indexer 完成后，
+提前为 Layer L+1 materialize prefix index K/scale。
+
 Layer L+1 attention 前，
 复用已经 materialize 完成的 prefix dense KV，
 只同步补齐 current/suffix pages。
+
+Layer L+1 topk/indexer 前，
+复用已经 materialize 完成的 prefix dense index buffer，
+只同步补齐 current/suffix pages。
 ```
 
 ### 3.2 性能目标
@@ -225,7 +233,6 @@ next layer consume waits only if prefetch has not finished
 
 Phase 8 不做：
 
-- 不做 NSA index K/scale prefetch；
 - 不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND`；
 - 不做 bandwidth throttle/page budget；
 - 不做多层 dense KV 常驻缓存；
@@ -576,7 +583,7 @@ disable MLA prefetch and fallback sync materialize
 2. 太早会和当前层 index materialize/topk 抢带宽；
 3. 太早会增加 NCCL collective 顺序风险。
 
-推荐在 `NativeSparseAttnBackend.forward_extend(...)` 中启动：
+MLA KV 推荐在 `NativeSparseAttnBackend.forward_extend(...)` 中启动：
 
 ```text
 1. 当前层需要的 kv_cache/page_table_1 已经准备好；
@@ -614,6 +621,19 @@ if prefetcher is not None:
 return attn_output
 ```
 
+Index K/scale 的启动点在 `Indexer.forward_cuda(...)` 末尾：
+
+```python
+topk_result = run_current_layer_topk(...)
+index_prefetcher.start_next_layer_prefix(
+    next_layer_id=layer_id + 1,
+    token_to_kv_pool=forward_batch.token_to_kv_pool,
+)
+return topk_result
+```
+
+当前层 `_maybe_materialize_shared_index_buffer(...)` 优先消费 `cp_shared_kv_index_prefetcher.consume(...)`；miss 时回退原来的 `materialize_shared_paged_buffer(...)`，因此关闭环境变量或 prefetch 条件不满足时语义不变。
+
 `wait_attention_window()` 仍在 `forward_extend(...)` 的 finally 中调用，但当前实现不做同步等待，只保留 deferred 日志/状态检查入口。真正的同步点在下一层 `consume(...)`：
 
 ```text
@@ -790,9 +810,10 @@ Phase 8 的 local copy/remap helper 应尽量复用 Phase 7 的 tai materialize
 ```text
 SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)
 forward_batch.cp_shared_kv_mla_prefetcher
+forward_batch.cp_shared_kv_index_prefetcher
 ```
 
-### Step 2: 拆分 token materialize helper
+### Step 2: 拆分 token/page materialize helper
 
 修改：
 
@@ -803,6 +824,8 @@ forward_batch.cp_shared_kv_mla_prefetcher
 ```text
 slot remap builder
 range local token KV materialize into existing dense buffer
+range local paged buffer materialize into existing dense buffer
+logical page -> slot dense page remap
 sync range all-reduce
 async all-reduce wrapper
 ```
@@ -813,7 +836,7 @@ async all-reduce wrapper
 full materialize 结果 == prefix materialize + suffix materialize 结果
 ```
 
-### Step 3: 新增 MLA prefetcher
+### Step 3: 新增 MLA / index prefetcher
 
 新增：
 
@@ -822,7 +845,7 @@ full materialize 结果 == prefix materialize + suffix materialize 结果
 职责：
 
 ```text
-gate
+common gate
 allocate dense buffer
 prefetch prefix range
 record event
@@ -831,7 +854,15 @@ defer wait until consume
 fallback
 ```
 
-### Step 4: 接入 `nsa_backend.forward_extend`
+Index K/scale prefetch 与 MLA prefetch 使用相同 slot-layout page table：
+
+```text
+real_page_table.reshape(-1)[slot] -> dense page slot + 1
+```
+
+区别是 MLA KV materialize 的单位是 token rows，需要 `page_size` 展开；index K/scale materialize 的单位是 page rows，prefix/suffix range 直接对应 dense page 行。
+
+### Step 4: 接入 `nsa_backend.forward_extend` / `nsa_indexer`
 
 修改：
 
@@ -840,7 +871,7 @@ fallback
 接入位置：
 
 ```text
-shared KV PAGED path 的 MLA materialize 分支。
+shared KV PAGED path 的 MLA materialize 分支，以及 NSA indexer 的 `_maybe_materialize_shared_index_buffer(...)`。
 ```
 
 逻辑：
@@ -853,6 +884,15 @@ run attention
 do not wait before return; consume waits before use
 ```
 
+Index K/scale 的启动点在当前层 topk/indexer 完成后：
+
+```text
+Layer L writes current layer index K/scale
+Layer L materializes/uses current layer index K/scale for topk
+Layer L starts Layer L+1 prefix index K/scale prefetch
+Layer L+1 indexer consumes prefetched prefix and fills suffix
+```
+
 ### Step 5: 单元测试
 
 修改：
@@ -861,12 +901,13 @@ do not wait before return; consume waits before use
 
 新增覆盖：
 
-1. prefix/suffix range materialize 拼接结果等价于 full materialize；
+1. token/paged prefix/suffix range materialize 拼接结果等价于 full materialize；
 2. prefix page-aligned gate；
 3. prefix_len=0 不 prefetch；
 4. non-PAGED / batch_size>1 / debug enabled fallback；
 5. consume handle layer mismatch fallback；
 6. started async handle 必须 wait。
+7. index materialize 优先消费 prefetched buffer，miss 后回退 full materialize；
 
 ### Step 6: 远端集成验证
 
@@ -885,7 +926,7 @@ router:  g0034
 3. 重复请求触发 radix hit；
 4. 检查输出质量；
 5. 检查没有 collective hang；
-6. profile 确认 `cp_shared_kv.materialize.token` 同步时间下降或被后续计算 overlap，并能看到下一层 consume 前的必要 wait。
+6. profile 确认 `cp_shared_kv.materialize.token` / index paged materialize 同步时间下降或被后续计算 overlap，并能看到下一层 consume 前的必要 wait。
 
 ---
 
@@ -972,4 +1013,3 @@ Fallback：
 debug enabled 时自动回到现有同步 materialize
 pynccl unavailable 时自动回到现有同步 materialize
 ```
-

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

@@ -10,13 +10,19 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     _all_reduce_materialized_buffer_async,
     _all_reduce_materialized_buffer_range,
     build_shared_token_kv_slot_remap,
+    build_slot_page_inverse_optimized,
+    build_slot_page_remap,
     cp_shared_kv_debug_enabled,
     cp_shared_kv_mla_prefetch_enabled,
     cp_shared_kv_mla_prefetch_log,
     cp_shared_kv_mla_prefetch_should_log_layer,
     filter_locs_mappable_to_physical_pool,
+    filter_pages_mappable_to_physical_pool,
+    materialize_local_paged_buffer_page_slots_into,
     materialize_local_token_kv_page_slots_into,
+    remap_logical_pages_to_slot_dense_pages,
     remap_logical_locs_to_slot_dense_locs_optimized,
+    slot_range_to_page_slice,
     slot_range_to_token_slice,
 )
 from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split
@@ -30,6 +36,14 @@ def _prefetch_log(message: str, *args) -> None:
     cp_shared_kv_mla_prefetch_log(message, *args)
 
 
+def _index_prefetch_fallback_log(reason: str, message: str, *args) -> None:
+    logger.info(
+        "CP shared KV index prefetch fallback (%s): " + message,
+        reason,
+        *args,
+    )
+
+
 def _is_cuda_stream_capturing() -> bool:
     if not torch.cuda.is_available():
         return False
@@ -46,6 +60,13 @@ class CpSharedKVMlaPrefetchHandle:
     event: torch.cuda.Event
 
 
+@dataclass
+class CpSharedKVIndexPrefetchHandle:
+    layer_id: int
+    dense_page_buffer: torch.Tensor
+    event: torch.cuda.Event
+
+
 class CpSharedKVMlaPrefetcher:
     """One-layer-ahead MLA prefix materialize prefetch for CP shared KV.
 
@@ -416,3 +437,423 @@ class CpSharedKVMlaPrefetcher:
             self.layout.cp_size,
             *args,
         )
+
+
+class CpSharedKVIndexPrefetcher:
+    """One-layer-ahead NSA index K/scale prefix materialize prefetch.
+
+    The index buffer is page-granular, unlike MLA KV which is token-granular.
+    It still uses the same slot-layout page table as MLA prefetch so the next
+    layer can consume the prefetched prefix and synchronously fill only current
+    suffix pages before running fp8 MQA/topk.
+    """
+
+    def __init__(
+        self,
+        *,
+        layout: CpSharedKVLayout,
+        prefix_pages: int,
+        slot_logical_pages: torch.Tensor,
+        page_inverse: torch.Tensor,
+        dense_num_pages: int,
+    ) -> None:
+        self.layout = layout
+        self.prefix_pages = prefix_pages
+        self.slot_logical_pages = slot_logical_pages
+        self.page_inverse = page_inverse
+        self.dense_num_pages = dense_num_pages
+        self.total_slots = int(slot_logical_pages.numel())
+        self.stream = torch.cuda.Stream()
+        self.handles: dict[int, CpSharedKVIndexPrefetchHandle] = {}
+        self.pending_attention_handle: Optional[CpSharedKVIndexPrefetchHandle] = None
+        self.disabled = False
+
+    @classmethod
+    def maybe_create(
+        cls,
+        *,
+        forward_batch: Any,
+        metadata: Any,
+        topk_transform_is_paged: bool,
+    ) -> Optional["CpSharedKVIndexPrefetcher"]:
+        if not cp_shared_kv_mla_prefetch_enabled():
+            return None
+        if cp_shared_kv_debug_enabled():
+            _index_prefetch_fallback_log(
+                "debug_enabled",
+                "SGLANG_DEBUG_CP_SHARED_KV is enabled.",
+            )
+            return None
+        if not torch.cuda.is_available() or _is_cuda_stream_capturing():
+            _index_prefetch_fallback_log(
+                "cuda_unavailable_or_stream_capturing",
+                "CUDA is unavailable or the current stream is capturing. cuda_available=%s",
+                torch.cuda.is_available(),
+            )
+            return None
+        if not getattr(forward_batch, "uses_cp_shared_kv", False):
+            _index_prefetch_fallback_log(
+                "not_cp_shared_kv",
+                "forward batch is not using CP shared KV.",
+            )
+            return None
+        if getattr(forward_batch, "hisparse_coordinator", None) is not None:
+            _index_prefetch_fallback_log(
+                "hisparse",
+                "HiSparse coordinator is active.",
+            )
+            return None
+        forward_mode = getattr(forward_batch, "forward_mode", None)
+        if forward_mode is None or not forward_mode.is_context_parallel_extend():
+            _index_prefetch_fallback_log(
+                "not_context_parallel_extend",
+                "forward mode is not context-parallel extend. forward_mode=%s",
+                forward_mode,
+            )
+            return None
+        if not is_nsa_prefill_cp_in_seq_split():
+            _index_prefetch_fallback_log(
+                "not_in_seq_split",
+                "NSA prefill CP mode is not in-seq-split.",
+            )
+            return None
+        if not topk_transform_is_paged:
+            _index_prefetch_fallback_log(
+                "not_paged_topk",
+                "topk transform is not PAGED.",
+            )
+            return None
+        if int(getattr(forward_batch, "batch_size", 0)) != 1:
+            _index_prefetch_fallback_log(
+                "batch_size",
+                "batch size is not supported. batch_size=%s",
+                getattr(forward_batch, "batch_size", None),
+            )
+            return None
+
+        token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None)
+        if token_to_kv_pool is None:
+            _index_prefetch_fallback_log(
+                "missing_token_to_kv_pool",
+                "forward batch has no token_to_kv_pool.",
+            )
+            return None
+        if getattr(token_to_kv_pool, "layer_transfer_counter", None) is not None:
+            _index_prefetch_fallback_log(
+                "layer_transfer_active",
+                "layer transfer is active.",
+            )
+            return None
+
+        layout = getattr(forward_batch, "cp_shared_kv_layout", None)
+        if layout is None:
+            _index_prefetch_fallback_log(
+                "missing_layout",
+                "forward batch has no CP shared KV layout.",
+            )
+            return None
+
+        extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None)
+        if extend_prefix_lens_cpu is None or len(extend_prefix_lens_cpu) != 1:
+            _index_prefetch_fallback_log(
+                "bad_prefix_lens_metadata",
+                "extend_prefix_lens_cpu is missing or not single-batch. value=%s",
+                extend_prefix_lens_cpu,
+            )
+            return None
+        page_size = int(getattr(token_to_kv_pool, "page_size", 1))
+        if page_size <= 1:
+            _index_prefetch_fallback_log(
+                "bad_page_size",
+                "page size is not supported. page_size=%s",
+                page_size,
+            )
+            return None
+        extend_prefix_len = int(extend_prefix_lens_cpu[0])
+        if extend_prefix_len <= 0 or extend_prefix_len % page_size != 0:
+            _index_prefetch_fallback_log(
+                "prefix_not_page_aligned",
+                "prefix length is zero or not page-aligned. prefix_len=%s page_size=%s",
+                extend_prefix_len,
+                page_size,
+            )
+            return None
+        prefix_pages = extend_prefix_len // page_size
+
+        real_page_table = getattr(metadata, "real_page_table", None)
+        page_table_1 = getattr(metadata, "page_table_1", None)
+        if real_page_table is None or page_table_1 is None:
+            _index_prefetch_fallback_log(
+                "missing_page_tables",
+                "metadata is missing real_page_table or page_table_1.",
+            )
+            return None
+        if prefix_pages <= 0 or prefix_pages > int(real_page_table.numel()):
+            _index_prefetch_fallback_log(
+                "prefix_pages_out_of_range",
+                "prefix pages are outside real page table. prefix_pages=%s real_pages=%s",
+                prefix_pages,
+                int(real_page_table.numel()),
+            )
+            return None
+
+        cp_group = get_attention_cp_group()
+        if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1:
+            _index_prefetch_fallback_log(
+                "missing_pynccl",
+                "pynccl communicator is unavailable. cp_rank=%s cp_size=%s",
+                layout.cp_rank,
+                layout.cp_size,
+            )
+            return None
+
+        try:
+            first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0))
+            page_buffer = token_to_kv_pool.get_index_k_with_scale_buffer(
+                layer_id=first_layer_id
+            )
+            remap_logical_pages = filter_pages_mappable_to_physical_pool(
+                logical_pages=real_page_table,
+                layout=layout,
+                physical_page_capacity=page_buffer.shape[0],
+            )
+            slot_logical_pages, _ = build_slot_page_remap(remap_logical_pages)
+            logical_page_capacity = max(int(page_buffer.shape[0]) - 1, 0) * (
+                layout.cp_size
+            ) + 1
+            page_inverse = build_slot_page_inverse_optimized(
+                slot_logical_pages,
+                logical_page_capacity=logical_page_capacity,
+            )
+        except Exception as exc:
+            _index_prefetch_fallback_log(
+                "init_exception",
+                "failed to initialize index prefetcher; falling back to sync materialize. error=%s",
+                exc,
+            )
+            logger.exception("Failed to initialize CP shared KV index prefetcher.")
+            return None
+
+        _prefetch_log(
+            "index_create cp_rank=%s cp_size=%s prefix_pages=%s total_slots=%s dense_pages=%s page_size=%s",
+            layout.cp_rank,
+            layout.cp_size,
+            prefix_pages,
+            int(slot_logical_pages.numel()),
+            int(slot_logical_pages.numel()) + 1,
+            page_size,
+        )
+
+        return cls(
+            layout=layout,
+            prefix_pages=prefix_pages,
+            slot_logical_pages=slot_logical_pages,
+            page_inverse=page_inverse,
+            dense_num_pages=int(slot_logical_pages.numel()) + 1,
+        )
+
+    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)
+        if kv_buffer is None:
+            return layer_id >= start_layer
+        return start_layer <= layer_id < start_layer + len(kv_buffer)
+
+    def consume(
+        self,
+        *,
+        layer_id: int,
+        page_buffer: torch.Tensor,
+        logical_pages: torch.Tensor,
+    ) -> Optional[tuple[torch.Tensor, torch.Tensor]]:
+        if self.disabled:
+            self._log_layer(
+                layer_id,
+                "index_consume_skip reason=disabled layer=%s",
+                layer_id,
+            )
+            return None
+
+        handle = self.handles.pop(layer_id, None)
+        if handle is None:
+            self._log_layer(layer_id, "index_consume_miss layer=%s", layer_id)
+            return None
+        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_skip reason=layer_mismatch expected=%s actual=%s",
+                layer_id,
+                handle.layer_id,
+            )
+            return None
+
+        torch.cuda.current_stream().wait_event(handle.event)
+        dense_page_buffer = handle.dense_page_buffer
+        suffix_slots = self.total_slots - self.prefix_pages
+
+        if self.prefix_pages < self.total_slots:
+            materialize_local_paged_buffer_page_slots_into(
+                page_buffer=page_buffer,
+                dense_page_buffer=dense_page_buffer,
+                slot_logical_pages=self.slot_logical_pages,
+                layout=self.layout,
+                start_slot=self.prefix_pages,
+                end_slot=self.total_slots,
+            )
+            suffix_rows = slot_range_to_page_slice(
+                self.prefix_pages,
+                self.total_slots,
+            )
+            _all_reduce_materialized_buffer_range(
+                dense_page_buffer,
+                self.layout.cp_size,
+                suffix_rows.start,
+                suffix_rows.stop,
+            )
+
+        self._log_layer(
+            layer_id,
+            "index_consume_hit layer=%s prefix_pages=%s suffix_slots=%s dense_pages=%s",
+            layer_id,
+            self.prefix_pages,
+            suffix_slots,
+            int(dense_page_buffer.shape[0]),
+        )
+
+        logical_pages = filter_pages_mappable_to_physical_pool(
+            logical_pages=logical_pages,
+            layout=self.layout,
+            physical_page_capacity=page_buffer.shape[0],
+        )
+        dense_pages = remap_logical_pages_to_slot_dense_pages(
+            logical_pages,
+            page_inverse=self.page_inverse,
+        )
+        return dense_page_buffer, dense_pages
+
+    def start_next_layer_prefix(
+        self,
+        *,
+        next_layer_id: int,
+        token_to_kv_pool: Any,
+    ) -> None:
+        if self.disabled:
+            self._log_next_layer(
+                next_layer_id,
+                "index_start_skip reason=disabled next_layer=%s",
+                next_layer_id,
+            )
+            return
+        if next_layer_id in self.handles:
+            self._log_next_layer(
+                next_layer_id,
+                "index_start_skip reason=already_started next_layer=%s",
+                next_layer_id,
+            )
+            return
+        if not self._layer_in_pool(token_to_kv_pool, next_layer_id):
+            self._log_next_layer(
+                next_layer_id,
+                "index_start_skip reason=layer_out_of_pool next_layer=%s",
+                next_layer_id,
+            )
+            return
+
+        try:
+            page_buffer = token_to_kv_pool.get_index_k_with_scale_buffer(
+                layer_id=next_layer_id
+            )
+        except Exception:
+            logger.exception(
+                "Failed to get next-layer index buffer for CP shared KV index prefetch."
+            )
+            self.disabled = True
+            self._log_next_layer(
+                next_layer_id,
+                "index_start_disable reason=get_index_buffer_failed next_layer=%s",
+                next_layer_id,
+            )
+            return
+
+        current_stream = torch.cuda.current_stream()
+        self.stream.wait_stream(current_stream)
+        try:
+            with torch.cuda.stream(self.stream):
+                dense_page_buffer = page_buffer.new_zeros(
+                    (self.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=self.slot_logical_pages,
+                    layout=self.layout,
+                    start_slot=0,
+                    end_slot=self.prefix_pages,
+                )
+                prefix_rows = slot_range_to_page_slice(0, self.prefix_pages)
+                event = _all_reduce_materialized_buffer_async(
+                    dense_page_buffer[prefix_rows],
+                    cp_size=self.layout.cp_size,
+                    stream=self.stream,
+                )
+                if event is None:
+                    self.disabled = True
+                    self._log_next_layer(
+                        next_layer_id,
+                        "index_start_disable reason=async_reduce_unavailable next_layer=%s",
+                        next_layer_id,
+                    )
+                    return
+        except Exception:
+            logger.exception("Failed to start CP shared KV index prefix prefetch.")
+            self.disabled = True
+            self._log_next_layer(
+                next_layer_id,
+                "index_start_disable reason=start_exception next_layer=%s",
+                next_layer_id,
+            )
+            return
+
+        handle = CpSharedKVIndexPrefetchHandle(
+            layer_id=next_layer_id,
+            dense_page_buffer=dense_page_buffer,
+            event=event,
+        )
+        self.handles[next_layer_id] = handle
+        self.pending_attention_handle = handle
+        self._log_next_layer(
+            next_layer_id,
+            "index_start next_layer=%s prefix_pages=%s dense_pages=%s",
+            next_layer_id,
+            self.prefix_pages,
+            int(dense_page_buffer.shape[0]),
+        )
+
+    def wait_attention_window(self) -> None:
+        handle = self.pending_attention_handle
+        if handle is None:
+            return
+        self._log_next_layer(
+            handle.layer_id,
+            "index_attention_wait_deferred next_layer=%s",
+            handle.layer_id,
+        )
+
+    def _log_layer(self, layer_id: int, message: str, *args) -> None:
+        if cp_shared_kv_mla_prefetch_should_log_layer(layer_id):
+            self._log(message, *args)
+
+    def _log_next_layer(self, next_layer_id: int, message: str, *args) -> None:
+        if cp_shared_kv_mla_prefetch_should_log_layer(next_layer_id):
+            self._log(message, *args)
+
+    def _log(self, message: str, *args) -> None:
+        _prefetch_log(
+            "cp_rank=%s cp_size=%s " + message,
+            self.layout.cp_rank,
+            self.layout.cp_size,
+            *args,
+        )

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

@@ -773,6 +773,43 @@ def remap_logical_locs_to_slot_dense_locs(
     return torch.where(mapped, dense_values, dense_locs)
 
 
+def remap_logical_pages_to_slot_dense_pages(
+    logical_pages: torch.Tensor,
+    page_inverse: torch.Tensor,
+) -> torch.Tensor:
+    """Map logical page ids through a fixed-shape slot page inverse.
+
+    This is the paged-buffer equivalent of
+    `remap_logical_locs_to_slot_dense_locs`.  Page 0 remains the dummy dense page
+    0; negative sentinels and pages absent from `page_inverse` remain `-1`.
+    """
+
+    dense_pages_out = torch.full_like(logical_pages, -1)
+    if logical_pages.numel() == 0 or page_inverse.numel() == 0:
+        return dense_pages_out
+
+    pages_long = logical_pages.to(torch.long)
+    valid_pages = pages_long >= 0
+    safe_pages = torch.where(valid_pages, pages_long, torch.zeros_like(pages_long))
+    pages_in_range = safe_pages < page_inverse.numel()
+    clamped_pages = torch.clamp(safe_pages, max=page_inverse.numel() - 1)
+    dense_pages = page_inverse[clamped_pages]
+    mapped = valid_pages & pages_in_range & (dense_pages >= 0)
+
+    if cp_shared_kv_debug_enabled() and torch.any(
+        valid_pages & pages_in_range & (dense_pages < 0)
+    ):
+        missing_pages = pages_long[valid_pages & pages_in_range & (dense_pages < 0)]
+        raise RuntimeError(
+            "CP shared KV slot remap got logical pages outside remap_logical_pages. "
+            f"missing_page_min={int(missing_pages.min().item())} "
+            f"missing_page_max={int(missing_pages.max().item())} "
+            f"logical_pages={tensor_debug_summary(logical_pages)}"
+        )
+
+    return torch.where(mapped, dense_pages.to(logical_pages.dtype), dense_pages_out)
+
+
 def build_shared_token_kv_slot_remap(
     kv_cache: torch.Tensor,
     logical_locs: torch.Tensor | None,
@@ -1261,7 +1298,50 @@ def materialize_local_paged_buffer_page_slots(
     if slot_logical_pages.numel() == 0:
         return dense_page_buffer
 
-    logical_pages = slot_logical_pages.reshape(-1).to(torch.long)
+    materialize_local_paged_buffer_page_slots_into(
+        page_buffer=page_buffer,
+        dense_page_buffer=dense_page_buffer,
+        slot_logical_pages=slot_logical_pages,
+        layout=layout,
+        start_slot=0,
+        end_slot=int(slot_logical_pages.numel()),
+    )
+    return dense_page_buffer
+
+
+def materialize_local_paged_buffer_page_slots_into(
+    page_buffer: torch.Tensor,
+    dense_page_buffer: torch.Tensor,
+    slot_logical_pages: torch.Tensor,
+    layout: CpSharedKVLayout,
+    start_slot: int,
+    end_slot: int | None = None,
+) -> None:
+    """Materialize a slot range into an existing dense paged buffer.
+
+    Dense page 0 is the dummy page; slot `i` writes dense page `i + 1`.
+    This is used by layer prefetch: prefix pages are written asynchronously,
+    then only suffix pages are filled synchronously when the layer is reached.
+    """
+
+    flat_slot_logical_pages = slot_logical_pages.reshape(-1)
+    total_slots = int(flat_slot_logical_pages.numel())
+    if end_slot is None:
+        end_slot = total_slots
+    if start_slot < 0 or end_slot < start_slot or end_slot > total_slots:
+        raise ValueError(
+            "Invalid CP shared KV paged slot materialize range: "
+            f"start_slot={start_slot} end_slot={end_slot} total_slots={total_slots}"
+        )
+    if start_slot == end_slot:
+        return
+    if dense_page_buffer.shape[0] < total_slots + 1:
+        raise ValueError(
+            "CP shared KV dense paged buffer is too small for slot materialize: "
+            f"dense_pages={dense_page_buffer.shape[0]} expected_at_least={total_slots + 1}"
+        )
+
+    logical_pages = flat_slot_logical_pages[start_slot:end_slot].to(torch.long)
     owned_mask = layout.owned_pages_mask(logical_pages)
     physical_pages = layout.logical_pages_to_physical(logical_pages).to(torch.long)
     safe_physical_pages = torch.where(
@@ -1272,8 +1352,21 @@ def materialize_local_paged_buffer_page_slots(
     gathered = page_buffer[safe_physical_pages]
     owned_view = owned_mask.view(-1, *([1] * (gathered.ndim - 1)))
     zero = torch.zeros((), dtype=page_buffer.dtype, device=page_buffer.device)
-    dense_page_buffer[1:].copy_(torch.where(owned_view, gathered, zero))
-    return dense_page_buffer
+    dense_page_buffer[start_slot + 1 : end_slot + 1].copy_(
+        torch.where(owned_view, gathered, zero)
+    )
+
+
+def slot_range_to_page_slice(
+    start_slot: int,
+    end_slot: int,
+) -> slice:
+    if start_slot < 0 or end_slot < start_slot:
+        raise ValueError(
+            "Invalid CP shared KV slot page slice range: "
+            f"start_slot={start_slot} end_slot={end_slot}"
+        )
+    return slice(start_slot + 1, end_slot + 1)
 
 
 def paged_copy_debug_checksum(

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

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextlib
+import logging
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
@@ -34,6 +35,7 @@ _is_cuda = is_cuda()
 _is_hip = is_hip()
 _is_npu = is_npu()
 _is_fp8_fnuz = is_fp8_fnuz()
+logger = logging.getLogger(__name__)
 if _is_cuda:
     try:
         import deep_gemm
@@ -53,6 +55,7 @@ from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.attention.nsa.utils import (
     cp_all_gather_rerange_output,
     get_cp_shared_kv_local_out_cache_loc,
+    get_cp_shared_kv_local_physical_out_cache_loc,
     is_nsa_enable_prefill_cp,
     is_nsa_prefill_cp_in_seq_split,
     log_cp_shared_kv_direct_write_fallback,
@@ -102,6 +105,24 @@ def _compute_contiguous_valid_cp_query_count(
     return max(0, min(actual_seq_q, valid_count))
 
 
+def _log_cp_shared_kv_index_prefetch_fallback(
+    reason: str,
+    message: str,
+    *args,
+) -> None:
+    """Log every index prefetch fallback event.
+
+    Warmup and real requests can hit the same reason independently; do not
+    dedupe here or later real fallbacks become invisible during profiling.
+    """
+
+    logger.info(
+        "CP shared KV index prefetch fallback (%s): " + message,
+        reason,
+        *args,
+    )
+
+
 class BaseIndexerMetadata(ABC):
     @abstractmethod
     def get_seqlens_int32(self) -> torch.Tensor:
@@ -282,6 +303,39 @@ class Indexer(MultiPlatformOp):
 
         assert forward_batch.cp_shared_kv_layout is not None
         layout = forward_batch.cp_shared_kv_layout
+        index_prefetcher = getattr(
+            forward_batch, "cp_shared_kv_index_prefetcher", None
+        )
+        if index_prefetcher is not None:
+            prefetched = index_prefetcher.consume(
+                layer_id=layer_id,
+                page_buffer=index_buffer,
+                logical_pages=logical_page_table,
+            )
+            if prefetched is not None:
+                if cp_shared_kv_debug_enabled():
+                    cp_shared_kv_debug_log(
+                        "index_materialize_prefetch_hit",
+                        "NSA index materialize prefetch hit cp_rank=%s layer=%s dense_pages=%s buffer_ck=%s",
+                        layout.cp_rank,
+                        layer_id,
+                        tensor_debug_summary(prefetched[1]),
+                        tensor_debug_checksum(prefetched[0]),
+                    )
+                return prefetched
+            start_layer = int(getattr(forward_batch.token_to_kv_pool, "start_layer", 0))
+            if int(layer_id) > start_layer:
+                _log_cp_shared_kv_index_prefetch_fallback(
+                    "consume_miss",
+                    "prefetcher did not provide layer buffer; falling back to "
+                    "sync paged materialize. layer=%s start_layer=%s cp_rank=%s "
+                    "logical_page_table_shape=%s",
+                    layer_id,
+                    start_layer,
+                    layout.cp_rank,
+                    tuple(logical_page_table.shape),
+                )
+
         if cp_shared_kv_debug_enabled():
             cp_shared_kv_debug_log(
                 "index_materialize_call",
@@ -303,9 +357,24 @@ class Indexer(MultiPlatformOp):
                 layer_id,
                 tensor_debug_summary(dense_pages),
                 tensor_debug_checksum(materialized),
-            )
+        )
         return materialized, dense_pages
 
+    def _maybe_start_next_layer_index_prefetch(
+        self,
+        forward_batch: ForwardBatch,
+        layer_id: int,
+    ) -> None:
+        index_prefetcher = getattr(
+            forward_batch, "cp_shared_kv_index_prefetcher", None
+        )
+        if index_prefetcher is None:
+            return
+        index_prefetcher.start_next_layer_prefix(
+            next_layer_id=layer_id + 1,
+            token_to_kv_pool=forward_batch.token_to_kv_pool,
+        )
+
     def _filter_shared_index_write(
         self,
         forward_batch: ForwardBatch,
@@ -1355,12 +1424,9 @@ class Indexer(MultiPlatformOp):
         if local_out_loc.numel() == 0:
             return True
 
-        assert forward_batch.cp_shared_kv_layout is not None
-        physical_out_loc = (
-            forward_batch.cp_shared_kv_layout.logical_locs_to_physical(
-                local_out_loc
-            ).contiguous()
-        )
+        physical_out_loc = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)
+        if physical_out_loc is None:
+            return False
         self._store_index_k_cache(
             forward_batch=forward_batch,
             layer_id=layer_id,
@@ -1579,7 +1645,7 @@ class Indexer(MultiPlatformOp):
                     forward_batch.nsa_cp_metadata is not None
                     and is_nsa_prefill_cp_in_seq_split()
                 ):
-                    return self._get_topk_in_seq_cp_pair(
+                    topk_result = self._get_topk_in_seq_cp_pair(
                         forward_batch,
                         layer_id,
                         q_fp8,
@@ -1604,6 +1670,7 @@ class Indexer(MultiPlatformOp):
                 topk=self.index_topk,
                 layer_id=layer_id,
             )
+        self._maybe_start_next_layer_index_prefetch(forward_batch, layer_id)
         return topk_result
 
     def forward_npu(

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

@@ -167,6 +167,8 @@ class PageAlignedInSeqSplitInfo:
 @dataclass
 class NSAContextParallelMetadata:
     split_list: List[int] = None
+    split_list_tensor: torch.Tensor = None
+    split_prefix_tensor: torch.Tensor = None
     max_rank_len: List[int] = None
     zigzag_index: List[int] = None
     per_rank_actual_token: List[int] = None
@@ -529,6 +531,44 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
     return local_out_cache_loc
 
 
+def get_cp_shared_kv_local_physical_out_cache_loc(forward_batch: "ForwardBatch"):
+    """Return cached physical rows for this CP rank's shared-KV direct writes.
+
+    `get_cp_shared_kv_local_out_cache_loc` returns logical shared-KV locs because
+    radix/scheduler/PD metadata are expressed in the global logical address
+    space.  Persistent writes into this rank's compact physical pool need the
+    same locs remapped through `CpSharedKVLayout`.  The logical loc tensor is
+    batch-scoped, not layer-scoped, so cache the physical remap on ForwardBatch
+    and reuse it for both MLA KV and NSA index K/scale writes across layers.
+    """
+
+    cached = getattr(forward_batch, "cp_local_physical_out_cache_loc", None)
+    if cached is not None:
+        return cached
+
+    local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+    if local_out_cache_loc is None:
+        return None
+
+    if local_out_cache_loc.numel() == 0:
+        forward_batch.cp_local_physical_out_cache_loc = local_out_cache_loc
+        return local_out_cache_loc
+
+    layout = getattr(forward_batch, "cp_shared_kv_layout", None)
+    if layout is None:
+        log_cp_shared_kv_direct_write_fallback(
+            "missing_layout",
+            "cp_shared_kv_layout is missing while building physical out_cache_loc",
+        )
+        return None
+
+    physical_out_cache_loc = layout.logical_locs_to_physical(
+        local_out_cache_loc
+    ).contiguous()
+    forward_batch.cp_local_physical_out_cache_loc = physical_out_cache_loc
+    return physical_out_cache_loc
+
+
 def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
     if is_nsa_prefill_cp_round_robin_split():
         cp_size = get_attention_cp_size()
@@ -655,7 +695,7 @@ def nsa_use_prefill_cp(forward_batch, nsa_enable_prefill_cp=None):
         return False
 
 
-def cp_attn_tp_all_gather_reorganazied_into_tensor(
+def _cp_attn_tp_all_gather_padded_tensor(
     input_: torch.Tensor, total_len, attn_tp_size, forward_batch, stream_op
 ):
     """
@@ -683,11 +723,14 @@ def cp_attn_tp_all_gather_reorganazied_into_tensor(
     get_attention_cp_group().cp_all_gather_into_tensor_async(
         input_tensor_all, input_, stream_op
     )
-    # step3
+    return input_tensor_all
+
+
+def _trim_cp_rank_padding_after_all_gather(input_tensor_all, forward_batch):
     outputs_list_max = list(
         torch.split(input_tensor_all, forward_batch.nsa_cp_metadata.max_rank_len, dim=0)
     )
-    outputs = torch.cat(
+    return torch.cat(
         [
             outputs_list_max[index][:per_rank_len]
             for index, per_rank_len in enumerate(
@@ -696,9 +739,114 @@ def cp_attn_tp_all_gather_reorganazied_into_tensor(
         ],
         dim=0,
     )
+
+
+def cp_attn_tp_all_gather_reorganazied_into_tensor(
+    input_: torch.Tensor, total_len, attn_tp_size, forward_batch, stream_op
+):
+    input_tensor_all = _cp_attn_tp_all_gather_padded_tensor(
+        input_, total_len, attn_tp_size, forward_batch, stream_op
+    )
+    # step3
+    outputs = _trim_cp_rank_padding_after_all_gather(input_tensor_all, forward_batch)
     return outputs
 
 
+def _log_tai_in_seq_rerange_fallback(reason: str, message: str, *args) -> None:
+    logger.info(
+        "CP NSA in-seq all-gather rerange tai-kernel fallback (%s): " + message,
+        reason,
+        *args,
+    )
+
+
+def _try_tai_in_seq_all_gather_rerange(
+    input_tensor_all: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    *,
+    hidden_size: int,
+    cp_size: int,
+) -> torch.Tensor | None:
+    metadata = forward_batch.nsa_cp_metadata
+    split_lens = getattr(metadata, "split_list_tensor", None)
+    split_prefix = getattr(metadata, "split_prefix_tensor", None)
+    split_list = getattr(metadata, "split_list", None)
+    max_rank_len = getattr(metadata, "max_rank_len", None)
+    if split_lens is None or split_prefix is None:
+        _log_tai_in_seq_rerange_fallback(
+            "missing_metadata",
+            "split_list_tensor or split_prefix_tensor is missing",
+        )
+        return None
+    if split_list is None or max_rank_len is None:
+        _log_tai_in_seq_rerange_fallback(
+            "missing_cpu_metadata",
+            "split_list or max_rank_len is missing",
+        )
+        return None
+    if not input_tensor_all.is_cuda or input_tensor_all.dim() != 2:
+        return None
+
+    total_tokens = int(sum(split_list))
+    if total_tokens == 0:
+        return input_tensor_all.new_empty((0, hidden_size))
+
+    max_segment_len = int(max(split_list))
+    max_rank_token = int(max_rank_len[0]) if len(max_rank_len) > 0 else 0
+    if max_segment_len <= 0 or max_rank_token <= 0:
+        return None
+    if input_tensor_all.shape[0] < max_rank_token * cp_size:
+        return None
+
+    try:
+        from tai_kernel.nsa_prefill import in_seq_all_gather_rerange
+    except Exception as exc:
+        _log_tai_in_seq_rerange_fallback(
+            "import_failed",
+            "failed to import tai_kernel.nsa_prefill.in_seq_all_gather_rerange: %s",
+            exc,
+        )
+        return None
+
+    try:
+        return in_seq_all_gather_rerange(
+            input_tensor_all,
+            split_lens,
+            split_prefix,
+            total_tokens=total_tokens,
+            hidden_size=hidden_size,
+            max_segment_len=max_segment_len,
+            max_rank_token=max_rank_token,
+            cp_size=cp_size,
+        )
+    except Exception as exc:
+        _log_tai_in_seq_rerange_fallback(
+            "kernel_failed",
+            "tai-kernel in-seq rerange failed: %s",
+            exc,
+        )
+        return None
+
+
+def _torch_in_seq_all_gather_rerange(
+    input_tensor_all: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    hidden_size: int,
+) -> torch.Tensor:
+    output_tensor = _trim_cp_rank_padding_after_all_gather(
+        input_tensor_all, forward_batch
+    )
+    outputs_list = list(
+        torch.split(
+            output_tensor, forward_batch.nsa_cp_metadata.reverse_split_len, dim=0
+        )
+    )
+    output_tensor = torch.cat(
+        [outputs_list[i] for i in forward_batch.nsa_cp_metadata.cp_reverse_index], dim=0
+    )
+    return output_tensor.view(-1, hidden_size)
+
+
 def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream):
     """
     # for in-seq-split
@@ -746,23 +894,26 @@ def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream):
         return output_tensor
 
     bs_seq_len, hidden_size = input_tensor.shape
-    output_tensor = cp_attn_tp_all_gather_reorganazied_into_tensor(
+    input_tensor_all = _cp_attn_tp_all_gather_padded_tensor(
         input_tensor,
         forward_batch.nsa_cp_metadata.total_seq_lens,
         cp_size,
         forward_batch,
         stream,
     )
-    outputs_list = list(
-        torch.split(
-            output_tensor, forward_batch.nsa_cp_metadata.reverse_split_len, dim=0
-        )
+    output_tensor = _try_tai_in_seq_all_gather_rerange(
+        input_tensor_all,
+        forward_batch,
+        hidden_size=hidden_size,
+        cp_size=cp_size,
     )
-    output_tensor = torch.cat(
-        [outputs_list[i] for i in forward_batch.nsa_cp_metadata.cp_reverse_index], dim=0
+    if output_tensor is not None:
+        return output_tensor
+    return _torch_in_seq_all_gather_rerange(
+        input_tensor_all,
+        forward_batch,
+        hidden_size,
     )
-    output_tensor = output_tensor.view(-1, hidden_size)
-    return output_tensor
 
 
 def calculate_cp_seq_idx(cp_chunks_len, seqs_len):
@@ -908,6 +1059,7 @@ def prepare_input_dp_with_cp_dsa(
             )
         )
     prefix_sum_list = list(accumulate(split_list))
+    split_prefix_list = [0] + prefix_sum_list[:-1]
 
     # TODO Support multi-batch-cp-split, multi-batch-cp support has accuracy issues
     # cp_seq_index = calculate_cp_seq_idx(split_list[:], seqs_len[:])
@@ -932,6 +1084,10 @@ def prepare_input_dp_with_cp_dsa(
 
     nsa_cp_metadata = NSAContextParallelMetadata(
         split_list=split_list,
+        split_list_tensor=torch.tensor(split_list, device="cuda", dtype=torch.int32),
+        split_prefix_tensor=torch.tensor(
+            split_prefix_list, device="cuda", dtype=torch.int32
+        ),
         max_rank_len=max_rank_len,
         zigzag_index=zigzag_index,
         per_rank_actual_token=per_rank_actual_token,

python/sglang/srt/layers/attention/nsa_backend.py

@@ -10,6 +10,7 @@ from sglang.srt.configs.model_config import get_nsa_index_topk, is_deepseek_nsa
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import (
+    CpSharedKVIndexPrefetcher,
     CpSharedKVMlaPrefetcher,
 )
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
@@ -603,6 +604,7 @@ class NativeSparseAttnBackend(
         batch_size = forward_batch.batch_size
         device = forward_batch.seq_lens.device
         forward_batch.cp_shared_kv_mla_prefetcher = None
+        forward_batch.cp_shared_kv_index_prefetcher = None
 
         if forward_batch.forward_mode.is_target_verify():
             draft_token_num = self.speculative_num_draft_tokens
@@ -884,6 +886,15 @@ class NativeSparseAttnBackend(
                 ),
             )
         )
+        forward_batch.cp_shared_kv_index_prefetcher = (
+            CpSharedKVIndexPrefetcher.maybe_create(
+                forward_batch=forward_batch,
+                metadata=metadata,
+                topk_transform_is_paged=(
+                    topk_transform_method == TopkTransformMethod.PAGED
+                ),
+            )
+        )
 
     def _cal_indexer_k_start_end(
         self,
@@ -1825,6 +1836,11 @@ class NativeSparseAttnBackend(
         finally:
             if mla_prefetcher is not None:
                 mla_prefetcher.wait_attention_window()
+            index_prefetcher = getattr(
+                forward_batch, "cp_shared_kv_index_prefetcher", None
+            )
+            if index_prefetcher is not None:
+                index_prefetcher.wait_attention_window()
 
         return attn_output
 

python/sglang/srt/model_executor/forward_batch_info.py

@@ -423,8 +423,10 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     uses_cp_shared_kv: bool = False
     cp_shared_kv_layout: Optional[CpSharedKVLayout] = None
     cp_local_out_cache_loc: Optional[torch.Tensor] = None
+    cp_local_physical_out_cache_loc: Optional[torch.Tensor] = None
     cp_shared_mla_direct_write_done: bool = False
     cp_shared_kv_mla_prefetcher: Optional[Any] = None
+    cp_shared_kv_index_prefetcher: Optional[Any] = None
 
     # For hidden states before normal
     return_hidden_states_before_norm: bool = False

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -8,6 +8,7 @@ from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cud
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.attention.nsa.utils import (
     get_cp_shared_kv_local_out_cache_loc,
+    get_cp_shared_kv_local_physical_out_cache_loc,
     log_cp_shared_kv_direct_write_fallback,
     nsa_use_prefill_cp,
 )
@@ -569,9 +570,11 @@ class DeepseekMLAForwardMixin:
         ):
             return True
 
-        physical_out_cache_loc = (
-            layout.logical_locs_to_physical(local_out_cache_loc).contiguous()
+        physical_out_cache_loc = get_cp_shared_kv_local_physical_out_cache_loc(
+            forward_batch
         )
+        if physical_out_cache_loc is None:
+            return False
         forward_batch.token_to_kv_pool.set_mla_kv_buffer(
             self.attn_mqa,
             physical_out_cache_loc,

test/registered/unit/layers/test_nsa_cp_utils.py

@@ -11,6 +11,7 @@ from sglang.srt.layers.attention.nsa.utils import (
     can_cp_split,
     cp_split_and_rebuild_1d,
     get_cp_shared_kv_local_out_cache_loc,
+    get_cp_shared_kv_local_physical_out_cache_loc,
     split_in_seq_cp_local_pair,
 )
 from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -335,6 +336,45 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             + list(range(6 * page_size, 7 * page_size)),
         )
 
+    def test_local_physical_out_cache_loc_is_cached(self):
+        import torch
+        from types import SimpleNamespace
+
+        page_size = 4
+        segment_pages = [1, 2, 3, 4, 8, 7, 6, 5]
+        out_cache_loc = torch.cat(
+            [
+                torch.arange(page * page_size, (page + 1) * page_size)
+                for page in segment_pages
+            ]
+        )
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=4,
+                cp_rank=1,
+            ),
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                split_list=[page_size] * 8,
+                zigzag_index=[1, 6],
+                page_aligned=True,
+                page_size=page_size,
+                extend_prefix_len=0,
+            ),
+            out_cache_loc=out_cache_loc,
+        )
+
+        physical_locs = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)
+        second_read = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)
+
+        self.assertIs(physical_locs, second_read)
+        self.assertEqual(
+            physical_locs.tolist(),
+            list(range(1 * page_size, 2 * page_size))
+            + list(range(2 * page_size, 3 * page_size)),
+        )
+
     def test_local_out_cache_loc_falls_back_when_owner_mismatch(self):
         import torch
         from types import SimpleNamespace

test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py

@@ -1214,6 +1214,254 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
         self.assertTrue(torch.equal(dense_page_buffer[1], page_buffer[1]))
         self.assertTrue(torch.equal(dense_page_buffer[2], page_buffer[2]))
 
+    def test_materialize_local_paged_buffer_page_slots_into_matches_full_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
+
+        layout = CpSharedKVLayout(page_size=64, cp_size=2, cp_rank=0)
+        page_buffer = torch.arange(0, 6 * 3, dtype=torch.float32).view(6, 3)
+        slot_logical_pages = torch.tensor([1, 2, 3, 4, 0, -1], dtype=torch.int64)
+
+        full = runtime.materialize_local_paged_buffer_page_slots(
+            page_buffer=page_buffer,
+            slot_logical_pages=slot_logical_pages,
+            layout=layout,
+        )
+        split = page_buffer.new_full(full.shape, -7)
+        split[0].zero_()
+
+        runtime.materialize_local_paged_buffer_page_slots_into(
+            page_buffer=page_buffer,
+            dense_page_buffer=split,
+            slot_logical_pages=slot_logical_pages,
+            layout=layout,
+            start_slot=0,
+            end_slot=3,
+        )
+        runtime.materialize_local_paged_buffer_page_slots_into(
+            page_buffer=page_buffer,
+            dense_page_buffer=split,
+            slot_logical_pages=slot_logical_pages,
+            layout=layout,
+            start_slot=3,
+            end_slot=slot_logical_pages.numel(),
+        )
+
+        self.assertTrue(torch.equal(split, full))
+
+    def test_remap_logical_pages_to_slot_dense_pages_preserves_sentinels(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+
+        slot_logical_pages = torch.tensor([1, 2, 0, 4, -1, 5], dtype=torch.int64)
+        page_inverse = runtime.build_slot_page_inverse(
+            slot_logical_pages,
+            logical_page_capacity=8,
+        )
+        logical_pages = torch.tensor([[0, 1, 5, -1, 3, 9]], dtype=torch.int32)
+
+        dense_pages = runtime.remap_logical_pages_to_slot_dense_pages(
+            logical_pages,
+            page_inverse=page_inverse,
+        )
+
+        self.assertEqual(dense_pages.tolist(), [[0, 1, 6, -1, -1, -1]])
+
+    def test_index_materialize_uses_prefetched_buffer_before_fallback(self):
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+
+        class FakePool:
+            def __init__(self):
+                self.index_buffer = torch.arange(0, 12, dtype=torch.float32).view(4, 3)
+
+            def get_index_k_with_scale_buffer(self, layer_id):
+                self.layer_id = layer_id
+                return self.index_buffer
+
+        class FakePrefetcher:
+            def __init__(self):
+                self.calls = []
+                self.dense_buffer = torch.full((3, 3), 5.0)
+                self.dense_pages = torch.tensor([[1, 2]], dtype=torch.int32)
+
+            def consume(self, *, layer_id, page_buffer, logical_pages):
+                self.calls.append((layer_id, page_buffer, logical_pages))
+                return self.dense_buffer, self.dense_pages
+
+        fake_pool = FakePool()
+        fake_prefetcher = FakePrefetcher()
+        forward_batch = SimpleNamespace(
+            token_to_kv_pool=fake_pool,
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=object(),
+            cp_shared_kv_index_prefetcher=fake_prefetcher,
+        )
+        logical_pages = torch.tensor([[1, 2]], dtype=torch.int32)
+        indexer = object.__new__(nsa_indexer.Indexer)
+
+        with patch.object(
+            nsa_indexer,
+            "materialize_shared_paged_buffer",
+            side_effect=AssertionError("prefetch hit must bypass full materialize"),
+        ):
+            dense_buffer, dense_pages = indexer._maybe_materialize_shared_index_buffer(
+                forward_batch,
+                layer_id=7,
+                logical_page_table=logical_pages,
+            )
+
+        self.assertIs(dense_buffer, fake_prefetcher.dense_buffer)
+        self.assertIs(dense_pages, fake_prefetcher.dense_pages)
+        self.assertEqual(fake_prefetcher.calls[0][0], 7)
+        self.assertIs(fake_prefetcher.calls[0][1], fake_pool.index_buffer)
+        self.assertIs(fake_prefetcher.calls[0][2], logical_pages)
+
+    def test_index_prefetch_start_targets_next_layer(self):
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+
+        class FakePrefetcher:
+            def __init__(self):
+                self.calls = []
+
+            def start_next_layer_prefix(self, *, next_layer_id, token_to_kv_pool):
+                self.calls.append((next_layer_id, token_to_kv_pool))
+
+        token_to_kv_pool = object()
+        fake_prefetcher = FakePrefetcher()
+        forward_batch = SimpleNamespace(
+            token_to_kv_pool=token_to_kv_pool,
+            cp_shared_kv_index_prefetcher=fake_prefetcher,
+        )
+        indexer = object.__new__(nsa_indexer.Indexer)
+
+        indexer._maybe_start_next_layer_index_prefetch(forward_batch, layer_id=11)
+
+        self.assertEqual(fake_prefetcher.calls, [(12, token_to_kv_pool)])
+
+    def test_index_prefetch_consume_miss_logs_fallback_after_first_layer(self):
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+
+        class FakePool:
+            start_layer = 0
+
+            def __init__(self):
+                self.index_buffer = torch.arange(0, 12, dtype=torch.float32).view(4, 3)
+
+            def get_index_k_with_scale_buffer(self, layer_id):
+                return self.index_buffer
+
+        class FakeLayout:
+            cp_rank = 3
+
+        class MissingPrefetcher:
+            def consume(self, *, layer_id, page_buffer, logical_pages):
+                return None
+
+        fallback_buffer = torch.full((3, 3), 9.0)
+        fallback_pages = torch.tensor([[1, 2]], dtype=torch.int32)
+        forward_batch = SimpleNamespace(
+            token_to_kv_pool=FakePool(),
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=FakeLayout(),
+            cp_shared_kv_index_prefetcher=MissingPrefetcher(),
+        )
+        logical_pages = torch.tensor([[1, 2]], dtype=torch.int32)
+        indexer = object.__new__(nsa_indexer.Indexer)
+
+        with patch.object(
+            nsa_indexer,
+            "materialize_shared_paged_buffer",
+            return_value=(fallback_buffer, fallback_pages),
+        ), patch(
+            "sglang.srt.layers.attention.nsa.nsa_indexer.logger",
+            create=True,
+        ) as logger:
+            dense_buffer, dense_pages = indexer._maybe_materialize_shared_index_buffer(
+                forward_batch,
+                layer_id=7,
+                logical_page_table=logical_pages,
+            )
+
+        self.assertIs(dense_buffer, fallback_buffer)
+        self.assertIs(dense_pages, fallback_pages)
+        logger.info.assert_called_once()
+        self.assertIn(
+            "CP shared KV index prefetch fallback",
+            logger.info.call_args.args[0],
+        )
+        self.assertIn("consume_miss", logger.info.call_args.args[1])
+
+    def test_index_prefetch_first_layer_miss_does_not_log_fallback(self):
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+
+        class FakePool:
+            start_layer = 0
+
+            def __init__(self):
+                self.index_buffer = torch.arange(0, 12, dtype=torch.float32).view(4, 3)
+
+            def get_index_k_with_scale_buffer(self, layer_id):
+                return self.index_buffer
+
+        class FakeLayout:
+            cp_rank = 0
+
+        class MissingPrefetcher:
+            def consume(self, *, layer_id, page_buffer, logical_pages):
+                return None
+
+        fallback_buffer = torch.full((3, 3), 9.0)
+        fallback_pages = torch.tensor([[1, 2]], dtype=torch.int32)
+        forward_batch = SimpleNamespace(
+            token_to_kv_pool=FakePool(),
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=FakeLayout(),
+            cp_shared_kv_index_prefetcher=MissingPrefetcher(),
+        )
+        logical_pages = torch.tensor([[1, 2]], dtype=torch.int32)
+        indexer = object.__new__(nsa_indexer.Indexer)
+
+        with patch.object(
+            nsa_indexer,
+            "materialize_shared_paged_buffer",
+            return_value=(fallback_buffer, fallback_pages),
+        ), patch(
+            "sglang.srt.layers.attention.nsa.nsa_indexer.logger",
+            create=True,
+        ) as logger:
+            indexer._maybe_materialize_shared_index_buffer(
+                forward_batch,
+                layer_id=0,
+                logical_page_table=logical_pages,
+            )
+
+        logger.info.assert_not_called()
+
+    def test_index_prefetch_create_skip_logs_fallback_when_enabled(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
+
+        with patch.object(
+            prefetch, "cp_shared_kv_mla_prefetch_enabled", return_value=True
+        ), patch.object(
+            prefetch, "cp_shared_kv_debug_enabled", return_value=False
+        ), patch.object(
+            prefetch.torch.cuda, "is_available", return_value=False
+        ), patch.object(
+            prefetch.logger, "info"
+        ) as logger:
+            result = prefetch.CpSharedKVIndexPrefetcher.maybe_create(
+                forward_batch=SimpleNamespace(),
+                metadata=SimpleNamespace(),
+                topk_transform_is_paged=True,
+            )
+
+        self.assertIsNone(result)
+        logger.assert_called_once()
+        self.assertIn(
+            "CP shared KV index prefetch fallback",
+            logger.call_args.args[0],
+        )
+        self.assertIn("cuda_unavailable_or_stream_capturing", logger.call_args.args[1])
+
 
 if __name__ == "__main__":
     unittest.main()