From 8844c303e23e998ccb1dc9e19da8f90e7a5c2ff6 Mon Sep 17 00:00:00 2001 From: ThomasX Date: Fri, 8 May 2026 03:20:47 +0800 Subject: [PATCH] fix: synchronize CP HiCache host eviction semantics --- .../sglang/srt/managers/cache_controller.py | 5 + python/sglang/srt/mem_cache/hiradix_cache.py | 38 ++++++- .../managers/test_hicache_controller_cp.py | 7 ++ .../mem_cache/test_cp_hicache_metadata.py | 100 +++++++++++++++++- 4 files changed, 144 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index a5959e9ee..e38c46c61 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -435,6 +435,11 @@ class HiCacheController: Requirement: no in-flight requests. This call is expected to run on the scheduler thread (control path), not concurrently with prefetch/backup. """ + if self.uses_cp_hicache: + raise RuntimeError( + "CP shared KV HiCache does not support attaching a storage backend at runtime." + ) + if self.enable_storage: raise RuntimeError("Storage backend already attached.") diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index aa4d1241f..3c27a7927 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -331,6 +331,12 @@ class HiRadixCache(RadixCache): prefetch/backup paths. Caller must ensure there are no running/queued requests to avoid races. """ + if self._uses_cp_hicache: + return ( + False, + "CP shared KV HiCache does not support attaching a storage backend at runtime.", + ) + # Validate inputs first (no side effects). if hicache_storage_prefetch_policy is not None: allowed = ["best_effort", "wait_complete", "timeout"] @@ -726,8 +732,14 @@ class HiRadixCache(RadixCache): device_indices=node.value, node_id=node.id, ) + required_host_slots = 0 + if getattr(result, "metadata", None) is None: + required_host_slots = result.required_host_slots + self._evict_host_for_physical_slots( + required_host_slots, + synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1, + ) if getattr(result, "metadata", None) is None: - self._evict_host_for_physical_slots(result.required_host_slots) result = self.cache_controller.write( device_indices=node.value, node_id=node.id, @@ -741,7 +753,7 @@ class HiRadixCache(RadixCache): self.ongoing_write_through[node.id] = node if not write_back: self.inc_node_lock_ref(node) - return len(node.cp_hicache.host_indices) + return len(node.value) host_indices = self.cache_controller.write( device_indices=node.value, @@ -1084,8 +1096,13 @@ class HiRadixCache(RadixCache): self._update_host_leaf_status(parent) return parent - def _evict_host_for_physical_slots(self, required_host_slots: int) -> int: - if required_host_slots <= 0: + def _evict_host_for_physical_slots( + self, required_host_slots: int, synchronize_across_ranks: bool = False + ) -> int: + synchronize_across_ranks = ( + synchronize_across_ranks and getattr(self, "tp_world_size", 1) > 1 + ) + if required_host_slots <= 0 and not synchronize_across_ranks: return 0 leaves = list(self.evictable_host_leaves) @@ -1095,7 +1112,18 @@ class HiRadixCache(RadixCache): heapq.heapify(eviction_heap) num_evicted = 0 - while num_evicted < required_host_slots and len(eviction_heap): + + def all_ranks_done() -> bool: + local_done = int(num_evicted >= required_host_slots) + if not synchronize_across_ranks: + return bool(local_done) + done = torch.tensor(local_done, dtype=torch.int, device="cpu") + torch.distributed.all_reduce( + done, op=torch.distributed.ReduceOp.MIN, group=self.tp_group + ) + return bool(done.item()) + + while len(eviction_heap) and not all_ranks_done(): _priority, x = heapq.heappop(eviction_heap) if x == self.root_node: break diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index cbaa1c381..02159f779 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -246,6 +246,13 @@ class TestHiCacheControllerCPWrite(CustomTestCase): self.assertEqual(config.model_name, "test-model") self.assertEqual(config.tp_lcm_size, 8) + def test_attach_storage_backend_rejects_cp_hicache(self): + host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) + controller = self.make_controller(host_pool) + + with self.assertRaisesRegex(RuntimeError, "CP shared KV.*storage backend"): + controller.attach_storage_backend("mooncake") + class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite): def test_cp_load_allocates_full_logical_locs_and_transfers_owned_physical_locs(self): diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 0f61eaa1d..f91f2da43 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -1,6 +1,6 @@ import sys import unittest -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import torch @@ -197,6 +197,19 @@ class FakeWriteController: return len(host_indices) +class FakeZeroOwnedWriteController: + write_policy = "write_through" + + def write(self, device_indices, node_id=-1, priority=None): + return FakeWriteSuccess( + CpHiCacheNodeMetadata( + logical_len=len(device_indices), + owned_positions=torch.empty((0,), dtype=torch.int64), + host_indices=torch.empty((0,), dtype=torch.int64), + ) + ) + + class FakeEvictionStrategy: def get_priority(self, node): return 0 @@ -286,6 +299,41 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertNotIn(1, root.children) self.assertEqual(node.host_len, 16) + def test_write_backup_cp_success_returns_logical_length_for_zero_owned_rank(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.cache_controller = FakeZeroOwnedWriteController() + cache.ongoing_write_through = {} + cache.inc_node_lock_ref = lambda node: None + node = TreeNode() + node.id = 123 + node.value = torch.arange(16, dtype=torch.int64) + + backed_len = cache.write_backup(node) + + self.assertEqual(backed_len, 16) + self.assertEqual(node.host_len, 16) + self.assertEqual(node.cp_hicache.host_indices.tolist(), []) + + def test_attach_storage_backend_rejects_cp_hicache_without_controller_call(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.cache_controller = type( + "Controller", + (), + { + "attach_storage_backend": lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("controller attach must not be called") + ) + }, + )() + + ok, message = cache.attach_storage_backend("mooncake") + + self.assertFalse(ok) + self.assertIn("CP shared KV", message) + self.assertIn("storage backend", message) + def test_evict_demotes_cp_backed_node_without_deleting_radix_child(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True @@ -515,6 +563,56 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): self.assertIsNotNone(parent.cp_hicache) self.assertEqual(parent.cp_hicache.host_indices.tolist(), [80]) + def test_synchronized_cp_host_eviction_removes_zero_owned_logical_leaf(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.tp_world_size = 2 + cache.tp_group = None + cache.root_node = TreeNode() + cache.root_node.key = RadixKey([]) + cache.evictable_host_leaves = set() + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache.eviction_strategy = type( + "Strategy", (), {"get_priority": lambda self, node: 0} + )() + cache._clear_pin = lambda node: None + cache._record_remove_event = lambda node: None + freed = [] + cache.cache_controller = type( + "Controller", + (), + { + "evict_host": lambda self, indices: freed.append(indices.clone()) + or len(indices) + }, + )() + node = TreeNode() + node.parent = cache.root_node + node.key = RadixKey([1, 2, 3, 4]) + node.value = None + 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), + ) + cache.root_node.children[1] = node + cache.evictable_host_leaves.add(node) + + def mark_not_done(done_tensor, op=None, group=None): + done_tensor.fill_(0) + + with patch("torch.distributed.all_reduce", side_effect=mark_not_done): + physical_freed = cache._evict_host_for_physical_slots( + 0, synchronize_across_ranks=True + ) + + self.assertEqual(physical_freed, 0) + self.assertEqual(freed, []) + self.assertNotIn(1, cache.root_node.children) + self.assertEqual(node.host_len, 0) + self.assertIsNone(node.cp_hicache) + class TestHiRadixCacheCPLoadBack(CustomTestCase): def test_cp_load_back_uses_host_len_not_host_value(self):