Stabilize CP HiCache page-first direct transfers on CUDA 13

CP HiCache direct/page_first_direct all-layer backup was still able to enter sgl-kernel's stale cudaMemcpyBatchAsync path, which segfaults under CUDA 13 before Python can surface an error. The SGLang route now avoids that all-layer sgl-kernel op for page_first_direct backup and uses the TAI per-layer direct LF->PF op for MHA, MLA, and NSA indexer data.\n\nThe load-back path also prepares NSA indexer page indices once per load op and reuses them across per-layer H2D loads, preserving per-layer overlap while removing redundant page-index derivation.\n\nConstraint: Remote runtime is CUDA 13.0 where sgl-kernel's all-layer direct LF->PF op uses the wrong cudaMemcpyBatchAsync ABI.\nRejected: Patch sgl-kernel in this branch | we are converging production HiCache direct/page_first_direct paths onto tai-kernel and do not want to maintain another CUDA-ABI-sensitive copy path here.\nRejected: Collapse H2D load-back into one all-layer op | that would reduce submit count but lose per-layer completion visibility and forward overlap.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce sgl_kernel.transfer_kv_all_layer_direct_lf_pf for direct/page_first_direct HiCache backup without CUDA 13 ABI verification.\nTested: g0034 container: PYTHONPATH=python python -m pytest -q -s test/registered/unit/mem_cache/test_nsa_pool_host_unit.py -> 10 passed, 3 warnings.\nTested: g0034 container: PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py -> 61 passed, 3 warnings.\nTested: python -m py_compile python/sglang/srt/mem_cache/memory_pool_host.py python/sglang/srt/managers/cache_controller.py test/registered/unit/mem_cache/test_nsa_pool_host_unit.py test/registered/unit/managers/test_hicache_controller_cp.py\nNot-tested: Full ETE prefill/decode traffic after this commit.\nNot-tested: sgl-kernel implementation itself remains unchanged.

test/registered/unit/managers/test_hicache_controller_cp.py

@@ -537,6 +537,61 @@ class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
         self.assertEqual(kwargs["layer_id"], 1)
         self.assertEqual(kwargs["page_size"], 1)
 
+    def test_nsa_indexer_load_reuses_precomputed_page_indices_across_layers(self):
+        calls = []
+
+        def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
+            calls.append(
+                (src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
+            )
+
+        host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
+        host_pool.layout = "page_first_direct"
+        host_pool.page_size = 4
+        host_pool.index_k_with_scale_buffer = torch.empty(
+            (8, 3, 1, 32), dtype=torch.uint8
+        )
+        device_pool = type("DevicePool", (), {})()
+        device_pool.index_k_with_scale_buffer = torch.empty(
+            (3, 8, 32), dtype=torch.uint8
+        )
+        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
+        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
+
+        original_getter = host_pool._get_indexer_page_indices
+        getter_calls = []
+
+        def counting_getter(h, d):
+            getter_calls.append((h.clone(), d.clone()))
+            return original_getter(h, d)
+
+        host_pool._get_indexer_page_indices = counting_getter
+
+        with patch(
+            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
+            return_value=fake_direct,
+        ):
+            host_pool.begin_load_to_device_op(
+                host_indices, device_indices, io_backend="direct"
+            )
+            try:
+                for layer_id in range(3):
+                    host_pool._load_indexer_to_device_per_layer(
+                        device_pool,
+                        host_indices,
+                        device_indices,
+                        layer_id=layer_id,
+                        io_backend="direct",
+                    )
+            finally:
+                host_pool.end_load_to_device_op()
+
+        self.assertEqual(len(getter_calls), 1)
+        self.assertEqual(len(calls), 3)
+        self.assertEqual([call[4]["layer_id"] for call in calls], [0, 1, 2])
+        self.assertEqual([call[2].tolist() for call in calls], [[1], [1], [1]])
+        self.assertEqual([call[3].tolist() for call in calls], [[3], [3], [3]])
+
 
 class FakeAllocator:
     def __init__(self, alloc_result=None):
@@ -1109,6 +1164,55 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
         self.assertEqual(host_pool.loads[0][0].tolist(), [100, 101, 102, 103])
         self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
 
+    def test_start_loading_prepares_load_op_once_for_all_layers(self):
+        class PreparingFakeHostPool(FakeHostPool):
+            def __init__(self, alloc_result):
+                super().__init__(alloc_result)
+                self.begin_calls = []
+                self.end_calls = 0
+                self.active = False
+
+            def begin_load_to_device_op(self, host_indices, device_indices, io_backend):
+                self.begin_calls.append(
+                    (host_indices.clone(), device_indices.clone(), io_backend)
+                )
+                self.active = True
+
+            def end_load_to_device_op(self):
+                self.end_calls += 1
+                self.active = False
+
+            def load_to_device_per_layer(
+                self, device_pool, host_indices, device_indices, layer_id, io_backend
+            ):
+                assert self.active
+                super().load_to_device_per_layer(
+                    device_pool, host_indices, device_indices, layer_id, io_backend
+                )
+
+        host_pool = PreparingFakeHostPool(
+            torch.tensor([100, 101, 102, 103], dtype=torch.int64)
+        )
+        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
+        allocator.device_pool = FakeDevicePool(layer_num=3)
+        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
+        node = TreeNode()
+        node.host_len = 16
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=16,
+            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
+            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
+            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
+            page_size=4,
+        )
+
+        controller.load_cp([node], node_id=113)
+        controller.start_loading()
+
+        self.assertEqual(len(host_pool.begin_calls), 1)
+        self.assertEqual(host_pool.end_calls, 1)
+        self.assertEqual([load[2] for load in host_pool.loads], [0, 1, 2])
+
     def test_cp_load_frees_unexpected_owner_allocator_length(self):
         host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
         allocator = FakeAllocator(alloc_result=torch.arange(64, 76, dtype=torch.int64))