feat: map CP HiCache loads to owned physical slots

python/sglang/srt/managers/cache_controller.py

@@ -790,6 +790,41 @@ class HiCacheController:
         )
         return device_indices
 
+    def load_cp(self, nodes_to_load, node_id: int = -1) -> Optional[torch.Tensor]:
+        logical_len = sum(node.host_len for node in nodes_to_load)
+        device_indices = self.mem_pool_device_allocator.alloc(logical_len)
+        if device_indices is None:
+            return None
+
+        host_chunks = []
+        physical_chunks = []
+        offset = 0
+        for node in nodes_to_load:
+            node_device_indices = device_indices[offset : offset + node.host_len]
+            offset += node.host_len
+            owned_positions = node.cp_hicache.owned_positions.to(device_indices.device)
+            if owned_positions.numel() == 0:
+                continue
+            selected_logical_locs = node_device_indices[owned_positions]
+            physical_chunks.append(
+                self.cp_shared_kv_layout.logical_locs_to_physical(selected_logical_locs)
+            )
+            host_chunks.append(node.cp_hicache.host_indices.to(device_indices.device))
+
+        if not host_chunks:
+            self._append_completed_load_ack(node_id)
+            return device_indices
+
+        self.load_queue.append(
+            CacheOperation(
+                torch.cat(host_chunks),
+                torch.cat(physical_chunks),
+                node_id,
+            )
+        )
+
+        return device_indices
+
     def move_indices(self, op: CacheOperation):
         host_indices, device_indices = op.host_indices, op.device_indices
         # move indices to GPU if using kernels, to host if using direct indexing

test/registered/unit/managers/test_hicache_controller_cp.py

@@ -29,6 +29,8 @@ if "sgl_kernel.kvcacheio" not in sys.modules:
 
 from sglang.srt.managers.cache_controller import HiCacheController
 from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
+from sglang.srt.mem_cache.radix_cache import TreeNode
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -206,5 +208,44 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
         self.assertEqual(result.required_host_slots, 4)
 
 
+class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
+    def test_cp_load_allocates_full_logical_locs_and_transfers_owned_physical_locs(self):
+        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
+        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
+        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),
+        )
+
+        device_indices = controller.load_cp([node], node_id=11)
+        controller.start_loading()
+
+        self.assertEqual(device_indices.tolist(), list(range(64, 80)))
+        self.assertEqual(allocator.alloc_calls, [16])
+        self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
+
+    def test_cp_load_zero_owned_returns_full_logical_locs_and_noop_ack(self):
+        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
+        allocator = FakeAllocator(alloc_result=torch.arange(64, 68, dtype=torch.int64))
+        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=3)
+        node = TreeNode()
+        node.host_len = 4
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=4,
+            owned_positions=torch.empty((0,), dtype=torch.int64),
+            host_indices=torch.empty((0,), dtype=torch.int64),
+        )
+
+        device_indices = controller.load_cp([node], node_id=12)
+
+        self.assertEqual(device_indices.tolist(), [64, 65, 66, 67])
+        self.assertEqual(host_pool.loads, [])
+        self.assertEqual(len(controller.ack_load_queue), 1)
+
+
 if __name__ == "__main__":
     main()