From f10a8983c3bf7aeaf57cbf21c5154046dd8adb3e Mon Sep 17 00:00:00 2001 From: ThomasX Date: Fri, 8 May 2026 02:26:46 +0800 Subject: [PATCH] fix: keep CP HiCache load host indices on CPU --- .../sglang/srt/managers/cache_controller.py | 2 +- .../managers/test_hicache_controller_cp.py | 32 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 235346708..ac69c3e73 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -809,7 +809,7 @@ class HiCacheController: 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)) + host_chunks.append(node.cp_hicache.host_indices) if not host_chunks: self._append_completed_load_ack(node_id) diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index a3b09b2e6..e122fa051 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -93,6 +93,15 @@ class FakeAllocator: return self.alloc_result[:need_size].clone() +class HostIndicesTensor(torch.Tensor): + @staticmethod + def __new__(cls, data): + return torch.Tensor._make_subclass(cls, data, require_grad=False) + + def to(self, *args, **kwargs): + raise AssertionError("load_cp should not move host indices before queuing") + + class DummyEvent: def record(self): pass @@ -246,6 +255,29 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite): self.assertEqual(host_pool.loads, []) self.assertEqual(len(controller.ack_load_queue), 1) + def test_cp_load_queues_cpu_host_indices_before_backend_moves(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) + host_indices = HostIndicesTensor(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) + node = TreeNode() + node.host_len = 16 + node.cp_hicache = type( + "CpHiCacheMetadataStub", + (), + { + "owned_positions": torch.tensor([4, 5, 6, 7], dtype=torch.int64), + "host_indices": host_indices, + }, + )() + + controller.load_cp([node], node_id=13) + + queued_op = controller.load_queue[0] + self.assertEqual(queued_op.host_indices.device.type, "cpu") + self.assertEqual(queued_op.host_indices.tolist(), [100, 101, 102, 103]) + self.assertEqual(queued_op.device_indices.tolist(), [20, 21, 22, 23]) + if __name__ == "__main__": main()