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.
2026-04-26 23:29:34 +08:00
parent 5af232e9de
commit d015e3fb01
5 changed files with 270 additions and 25 deletions
@@ -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,