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