feat: map CP HiCache loads to owned physical slots

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()