Batch CP shared-KV index work for bs>1 fast paths

The bs>1 path needs index top-k, shared-index prepare, current-index compact, and current-slot compose to consume flattened batch descriptors instead of falling back to per-request or per-segment Python/Torch work. This change wires SGLang to the new TAI batch prepare kernels, keeps fallback explicit, and records the remaining HiCache/load-backup gaps in the bs>1 workstream docs. Constraint: CP shared-KV bs>1 must reuse fast paths rather than adding slow batch-only fallbacks Constraint: No new collective operations were introduced Rejected: Leave current-only cp_index as Python slice/cat | it keeps per-segment overhead in the short-extend bs>1 case Rejected: Infer max segment lengths from CUDA descriptor tensors | .item() would add CPU synchronization on the hot path Confidence: medium Scope-risk: moderate Directive: Do not remove the explicit fallback warnings without verifying the corresponding TAI symbols are present in production Tested: local py_compile for touched SGLang files Tested: remote g0034 test_nsa_cp_utils.py passed, 53 tests Tested: remote g0034 test_fill_current_index_page_slots_uses_tai_kernel_when_available passed Not-tested: full ETE bs>1 traffic with HiCache load/backup and draft/EAGLE enabled
2026-06-03 08:45:32 +08:00
parent d36f62a3cd
commit 19dcd6c4dc
6 changed files with 1132 additions and 113 deletions
@@ -1191,6 +1191,55 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
        self.assertEqual(current_mask.tolist(), [[False, True, True, False, False]])
        self.assertEqual(mixed_locs.tolist(), [[4, 12, 13, -1, -1]])

+    def test_fill_current_index_page_slots_uses_tai_kernel_when_available(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+
+        class FakeKernels:
+            calls = []
+
+            @staticmethod
+            def fill_current_index_page_slots(*args, **kwargs):
+                FakeKernels.calls.append((args, kwargs))
+                dense_page_buffer = args[0]
+                dense_page_buffer[2, 0] = 99
+                return dense_page_buffer
+
+        dense_page_buffer = torch.zeros((3, 32), dtype=torch.uint8)
+        current_k = torch.ones((2, 4), dtype=torch.uint8)
+        current_scale = torch.ones((2, 1), dtype=torch.float32)
+        current_locs = torch.tensor([8, 9], dtype=torch.int64)
+        page_inverse = torch.tensor([0, -1, 2], dtype=torch.int64)
+
+        with patch.object(
+            runtime,
+            "_tai_materialize_runtime_enabled",
+            return_value=True,
+        ), patch.object(
+            runtime,
+            "_load_tai_materialize_kernels",
+            return_value=FakeKernels,
+        ):
+            result = runtime.fill_current_index_page_slots(
+                dense_page_buffer=dense_page_buffer,
+                current_index_k=current_k,
+                current_index_scale=current_scale,
+                current_locs=current_locs,
+                page_inverse=page_inverse,
+                page_size=4,
+                index_head_dim=4,
+            )
+
+        self.assertIs(result, dense_page_buffer)
+        self.assertEqual(int(result[2, 0]), 99)
+        self.assertEqual(len(FakeKernels.calls), 1)
+        args, kwargs = FakeKernels.calls[0]
+        self.assertIs(args[0], dense_page_buffer)
+        self.assertTrue(torch.equal(args[1], current_k))
+        self.assertTrue(torch.equal(args[2], current_scale))
+        self.assertTrue(torch.equal(args[3], current_locs))
+        self.assertTrue(torch.equal(args[4], page_inverse))
+        self.assertEqual(kwargs, {"page_size": 4, "index_head_dim": 4})
+
    def test_tai_current_slot_fill_is_skipped_when_sparse_page_self_test_fails(self):
        from sglang.srt.environ import envs
        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime