Reduce CP shared KV current-chunk materialization

Phase 2 shared KV keeps persistent KV/index sharded across CP ranks but
uses a full-view compatibility layer before NSA topk and MLA attention.
For current-only prefill chunks, the current KV/index tensors have already
been CP all-gathered and reranged before being written to the sharded
persistent pool. This change adds a guarded current-reuse path that remaps
logical current locs to compact tensor rows and skips the shared-KV
materialize path for current-only MLA and NSA indexer reads.

Constraint: Existing NSA/MLA kernels still consume full-view/compact page tables; history and mixed current/history batches must keep the Phase 2 fallback.
Rejected: Make all history attention shard-aware in this patch | that requires global topk merge and distributed sparse attention and belongs to a later phase.
Confidence: medium
Scope-risk: moderate
Reversibility: clean
Directive: Do not remove the Phase 2 fallback until mixed/history shared-KV paths have correctness and performance coverage.
Tested: python -m py_compile on modified Python files
Tested: git diff --check on staged modified files
Not-tested: local pytest collection is blocked by missing pybase64 in this environment.
Not-tested: full long-context chunked prefill/decode performance in this commit step.

python/sglang/srt/environ.py

@@ -203,6 +203,7 @@ class Envs:
     SGLANG_FORCE_SHUTDOWN = EnvBool(False)
     SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)
     SGLANG_DEBUG_CP_SHARED_KV = EnvBool(False)
+    SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)
     SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
     SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)
     SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)

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

@@ -17,6 +17,44 @@ def cp_shared_kv_debug_enabled() -> bool:
     return envs.SGLANG_DEBUG_CP_SHARED_KV.get()
 
 
+def cp_shared_kv_current_reuse_enabled() -> bool:
+    return envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.get()
+
+
+def is_current_only_extend_batch(forward_batch) -> bool:
+    """Return whether an extend batch has no cached/history tokens.
+
+    This intentionally uses CPU metadata instead of scanning CUDA page tables in
+    Python control flow.  It is a conservative gate for Phase 3 current reuse:
+    when it returns true, `seq_lens == extend_seq_lens` and all prefix lengths
+    are zero, so the logical KV view for the batch should be exactly the
+    current `out_cache_loc` chunk.
+    """
+
+    if forward_batch is None:
+        return False
+    forward_mode = getattr(forward_batch, "forward_mode", None)
+    if forward_mode is None or not forward_mode.is_extend_without_speculative():
+        return False
+
+    extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None)
+    extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
+    seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None)
+    if (
+        extend_prefix_lens_cpu is None
+        or extend_seq_lens_cpu is None
+        or seq_lens_cpu is None
+    ):
+        return False
+
+    if any(int(prefix_len) != 0 for prefix_len in extend_prefix_lens_cpu):
+        return False
+
+    seq_lens_list = [int(seq_len) for seq_len in seq_lens_cpu.tolist()]
+    extend_seq_lens_list = [int(seq_len) for seq_len in extend_seq_lens_cpu]
+    return seq_lens_list == extend_seq_lens_list
+
+
 def cp_shared_kv_debug_log(
     key: str,
     message: str,
@@ -230,6 +268,49 @@ def remap_logical_locs_to_dense_locs(
     return dense_locs
 
 
+def build_current_loc_remap(
+    query_locs: torch.Tensor,
+    current_locs: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Map logical locs into rows of the gathered current chunk tensor.
+
+    `query_locs` may be any shape and may contain `-1` invalid sentinels.
+    `current_locs` is `forward_batch.out_cache_loc`; its row order is the row
+    order of the already CP-all-gathered current KV/index tensor.
+
+    Returns:
+      - is_current_mask: true where query_locs is a non-negative current loc.
+      - compact_row_ids: row id into the current compact tensor where valid,
+        and -1 elsewhere.  The dtype/shape match query_locs.
+    """
+
+    is_current = torch.zeros_like(query_locs, dtype=torch.bool)
+    compact_row_ids = torch.full_like(query_locs, -1)
+    if query_locs.numel() == 0 or current_locs.numel() == 0:
+        return is_current, compact_row_ids
+
+    query_flat_long = query_locs.reshape(-1).to(torch.long)
+    current_flat_long = current_locs.reshape(-1).to(torch.long)
+    sorted_current_locs, sorted_to_current_rows = torch.sort(current_flat_long)
+
+    insert_positions = torch.searchsorted(sorted_current_locs, query_flat_long)
+    safe_positions = torch.clamp(insert_positions, max=current_flat_long.numel() - 1)
+    in_range = insert_positions < current_flat_long.numel()
+    matched = (
+        in_range
+        & (query_flat_long >= 0)
+        & (sorted_current_locs[safe_positions] == query_flat_long)
+    )
+
+    matched_rows = sorted_to_current_rows[safe_positions].to(compact_row_ids.dtype)
+    compact_flat = torch.where(
+        matched,
+        matched_rows,
+        torch.full_like(matched_rows, -1),
+    )
+    return matched.reshape(query_locs.shape), compact_flat.reshape(query_locs.shape)
+
+
 def logical_pages_from_locs(logical_locs: torch.Tensor, page_size: int) -> torch.Tensor:
     logical_pages = logical_locs.clone()
     valid_mask = logical_locs >= 0

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

@@ -16,7 +16,9 @@ from sglang.srt.layers.attention.nsa import index_buf_accessor
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     cp_shared_kv_debug_enabled,
     cp_shared_kv_debug_log,
+    cp_shared_kv_current_reuse_enabled,
     filter_owned_logical_locs,
+    is_current_only_extend_batch,
     materialize_shared_paged_buffer,
     tensor_debug_checksum,
     tensor_debug_summary,
@@ -319,6 +321,19 @@ class Indexer(MultiPlatformOp):
 
         return physical_out_loc, key[owned_mask].contiguous()
 
+    def _can_reuse_current_index_kv(self, forward_batch: ForwardBatch) -> bool:
+        return (
+            cp_shared_kv_current_reuse_enabled()
+            and forward_batch.uses_cp_shared_kv
+            and self.nsa_enable_prefill_cp
+            and forward_batch.nsa_cp_metadata is not None
+            and is_nsa_prefill_cp_in_seq_split()
+            and is_current_only_extend_batch(forward_batch)
+            and forward_batch.hisparse_coordinator is None
+            and _is_cuda
+            and not _is_fp8_fnuz
+        )
+
     @contextlib.contextmanager
     def _with_real_sm_count(self):
         # When pipeline parallelism is enabled, each PP rank initiates a recv operation after the _pp_launch_batch
@@ -822,6 +837,7 @@ class Indexer(MultiPlatformOp):
         kv_len: int,
         actual_seq_q: int,
         cp_index: List[Tuple[int, int, int]] = None,
+        current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ) -> torch.Tensor:
         if TYPE_CHECKING:
             assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool)
@@ -839,11 +855,29 @@ class Indexer(MultiPlatformOp):
         batch_idx_list = []
 
         block_tables = metadata.get_page_table_64()
-        index_buffer, block_tables = self._maybe_materialize_shared_index_buffer(
-            forward_batch,
-            layer_id,
-            block_tables,
-        )
+        if current_index_kv is not None and cp_index is not None:
+            current_index_kv = None
+        if current_index_kv is None:
+            index_buffer, block_tables = self._maybe_materialize_shared_index_buffer(
+                forward_batch,
+                layer_id,
+                block_tables,
+            )
+        else:
+            index_buffer = None
+            if cp_shared_kv_debug_enabled():
+                cp_shared_kv_debug_log(
+                    "index_current_reuse",
+                    "NSA index current reuse cp_rank=%s layer=%s kv_len=%s actual_seq_q=%s k_ck=%s s_ck=%s",
+                    forward_batch.cp_shared_kv_layout.cp_rank
+                    if forward_batch.cp_shared_kv_layout is not None
+                    else None,
+                    layer_id,
+                    kv_len,
+                    actual_seq_q,
+                    tensor_debug_checksum(current_index_kv[0]),
+                    tensor_debug_checksum(current_index_kv[1]),
+                )
 
         assert (
             forward_batch.seq_lens_cpu is not None
@@ -949,21 +983,27 @@ class Indexer(MultiPlatformOp):
             weights = weights[valid_q_mask].contiguous()
             ke_offset = ke_offset[valid_q_mask].contiguous()
             kv_len = min(cp_kv_end, logical_kv_limit)
-            k_fp8 = index_buf_accessor.GetK.execute(
-                forward_batch.token_to_kv_pool,
-                index_buffer,
-                seq_len=kv_len,
-                page_indices=block_tables[0],
-            )
-            k_scale = index_buf_accessor.GetS.execute(
-                forward_batch.token_to_kv_pool,
-                index_buffer,
-                seq_len=kv_len,
-                page_indices=block_tables[0],
-            )
+            if current_index_kv is None:
+                assert index_buffer is not None
+                k_fp8 = index_buf_accessor.GetK.execute(
+                    forward_batch.token_to_kv_pool,
+                    index_buffer,
+                    seq_len=kv_len,
+                    page_indices=block_tables[0],
+                )
+                k_scale = index_buf_accessor.GetS.execute(
+                    forward_batch.token_to_kv_pool,
+                    index_buffer,
+                    seq_len=kv_len,
+                    page_indices=block_tables[0],
+                )
 
-            k_fp8 = k_fp8.view(torch.float8_e4m3fn)
-            k_scale = k_scale.view(torch.float32).squeeze(-1)
+                k_fp8 = k_fp8.view(torch.float8_e4m3fn)
+                k_scale = k_scale.view(torch.float32).squeeze(-1)
+            else:
+                k_fp8, k_scale = current_index_kv
+                k_fp8 = k_fp8[:kv_len].contiguous()
+                k_scale = k_scale[:kv_len].view(torch.float32).squeeze(-1).contiguous()
             kv_fp8 = (k_fp8, k_scale)
             ks = torch.full((valid_q_count,), offset, dtype=torch.int32, device="cuda")
             ke = ks + ke_offset
@@ -1269,6 +1309,34 @@ class Indexer(MultiPlatformOp):
 
             weights = self._get_logits_head_gate(x_for_gate, q_scale)
 
+        current_index_kv = None
+        if self._can_reuse_current_index_kv(forward_batch):
+            if key.shape[0] == forward_batch.out_cache_loc.numel():
+                current_k_fp8, current_k_scale = act_quant(
+                    key.contiguous(), self.block_size, self.scale_fmt
+                )
+                current_index_kv = (
+                    current_k_fp8.contiguous(),
+                    current_k_scale.contiguous(),
+                )
+                if cp_shared_kv_debug_enabled():
+                    cp_shared_kv_debug_log(
+                        "index_current_reuse_prepare",
+                        "NSA index current reuse prepared cp_rank=%s layer=%s locs=%s k_ck=%s s_ck=%s",
+                        forward_batch.cp_shared_kv_layout.cp_rank
+                        if forward_batch.cp_shared_kv_layout is not None
+                        else None,
+                        layer_id,
+                        tensor_debug_summary(forward_batch.out_cache_loc),
+                        tensor_debug_checksum(current_index_kv[0]),
+                        tensor_debug_checksum(current_index_kv[1]),
+                    )
+            elif cp_shared_kv_debug_enabled():
+                raise RuntimeError(
+                    "CP shared KV current index reuse shape mismatch: "
+                    f"key_tokens={key.shape[0]} out_cache_loc={forward_batch.out_cache_loc.numel()}"
+                )
+
         if _is_cuda or _is_hip:
             assert forward_batch.seq_lens_cpu is not None
             if len(forward_batch.seq_lens_cpu) == 0:
@@ -1320,6 +1388,7 @@ class Indexer(MultiPlatformOp):
                         metadata,
                         kv_len_prev,
                         actual_seq_q_prev,
+                        current_index_kv=current_index_kv,
                     )
 
                     topk_result_next = self._get_topk_ragged_with_cp(
@@ -1330,6 +1399,7 @@ class Indexer(MultiPlatformOp):
                         metadata,
                         kv_len_next,
                         actual_seq_q_next,
+                        current_index_kv=current_index_kv,
                     )
                     return torch.cat([topk_result_prev, topk_result_next], dim=0)
                 else:

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

@@ -10,9 +10,12 @@ 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_runtime import (
+    build_current_loc_remap,
     cp_shared_kv_debug_enabled,
     cp_shared_kv_debug_log,
+    cp_shared_kv_current_reuse_enabled,
     filter_owned_logical_locs,
+    is_current_only_extend_batch,
     materialize_shared_token_kv_buffer,
     tensor_debug_checksum,
     tensor_debug_summary,
@@ -1591,13 +1594,49 @@ class NativeSparseAttnBackend(
             and topk_transform_method == TopkTransformMethod.PAGED
         ):
             assert forward_batch.cp_shared_kv_layout is not None
-            kv_cache, page_table_1 = materialize_shared_token_kv_buffer(
-                kv_cache=kv_cache,
-                logical_locs=page_table_1,
-                remap_logical_locs=metadata.page_table_1,
-                layout=forward_batch.cp_shared_kv_layout,
-                page_size=forward_batch.token_to_kv_pool.page_size,
+            can_reuse_current_kv = (
+                cp_shared_kv_current_reuse_enabled()
+                and is_current_only_extend_batch(forward_batch)
+                and k is not None
+                and k_rope is not None
+                and k.shape[0] == forward_batch.out_cache_loc.numel()
+                and k_rope.shape[0] == forward_batch.out_cache_loc.numel()
             )
+            if can_reuse_current_kv:
+                logical_page_table_1 = page_table_1
+                current_mask, page_table_1 = build_current_loc_remap(
+                    logical_page_table_1,
+                    forward_batch.out_cache_loc,
+                )
+                if cp_shared_kv_debug_enabled():
+                    missing_current = (logical_page_table_1 >= 0) & (~current_mask)
+                    if torch.any(missing_current):
+                        bad_locs = logical_page_table_1[missing_current]
+                        raise RuntimeError(
+                            "CP shared KV current MLA reuse expected current-only "
+                            "logical locs but found history locs. "
+                            f"bad_min={int(bad_locs.min().item())} "
+                            f"bad_max={int(bad_locs.max().item())}"
+                        )
+                    cp_shared_kv_debug_log(
+                        "mla_current_reuse",
+                        "MLA current reuse cp_rank=%s layer=%s current_locs=%s remapped=%s kv_ck=%s rope_ck=%s",
+                        forward_batch.cp_shared_kv_layout.cp_rank,
+                        layer.layer_id,
+                        tensor_debug_summary(forward_batch.out_cache_loc),
+                        tensor_debug_summary(page_table_1),
+                        tensor_debug_checksum(k),
+                        tensor_debug_checksum(k_rope),
+                    )
+                kv_cache = _cat([k, k_rope], dim=-1)
+            else:
+                kv_cache, page_table_1 = materialize_shared_token_kv_buffer(
+                    kv_cache=kv_cache,
+                    logical_locs=page_table_1,
+                    remap_logical_locs=metadata.page_table_1,
+                    layout=forward_batch.cp_shared_kv_layout,
+                    page_size=forward_batch.token_to_kv_pool.page_size,
+                )
 
         if nsa_impl == "tilelang":
             if q_rope is not None:

test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py

@@ -104,6 +104,60 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
 
         self.assertEqual(dense_locs.tolist(), [0, 1, 4, 5, -1, 8])
 
+    def test_build_current_loc_remap_supports_non_contiguous_locs_and_sentinel(self):
+        from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
+            build_current_loc_remap,
+        )
+
+        current_locs = torch.tensor([100, 64, 256, 128], dtype=torch.int64)
+        query_locs = torch.tensor(
+            [[128, -1, 64], [512, 100, 256]], dtype=torch.int32
+        )
+
+        is_current, compact_rows = build_current_loc_remap(query_locs, current_locs)
+
+        self.assertEqual(
+            is_current.tolist(),
+            [[True, False, True], [False, True, True]],
+        )
+        self.assertEqual(compact_rows.tolist(), [[3, -1, 1], [-1, 0, 2]])
+        self.assertEqual(compact_rows.dtype, query_locs.dtype)
+
+    def test_build_current_loc_remap_returns_all_invalid_for_empty_current_locs(self):
+        from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
+            build_current_loc_remap,
+        )
+
+        query_locs = torch.tensor([4, -1, 8], dtype=torch.int64)
+        is_current, compact_rows = build_current_loc_remap(
+            query_locs, torch.empty((0,), dtype=torch.int64)
+        )
+
+        self.assertEqual(is_current.tolist(), [False, False, False])
+        self.assertEqual(compact_rows.tolist(), [-1, -1, -1])
+
+    def test_is_current_only_extend_batch_uses_cpu_lengths_without_tensor_scans(self):
+        from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
+            is_current_only_extend_batch,
+        )
+        from sglang.srt.model_executor.forward_batch_info import ForwardMode
+
+        forward_batch = SimpleNamespace(
+            forward_mode=ForwardMode.EXTEND,
+            extend_prefix_lens_cpu=[0, 0],
+            extend_seq_lens_cpu=[3, 5],
+            seq_lens_cpu=torch.tensor([3, 5], dtype=torch.int32),
+        )
+
+        self.assertTrue(is_current_only_extend_batch(forward_batch))
+
+        forward_batch.extend_prefix_lens_cpu = [0, 1]
+        self.assertFalse(is_current_only_extend_batch(forward_batch))
+
+        forward_batch.extend_prefix_lens_cpu = [0, 0]
+        forward_batch.seq_lens_cpu = torch.tensor([4, 5], dtype=torch.int32)
+        self.assertFalse(is_current_only_extend_batch(forward_batch))
+
     def test_materialize_local_token_kv_pages(self):
         from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
             build_dense_page_remap,