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

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,