From bebbbfd0e8fa77e06f4cfc00d694bd9fad488492 Mon Sep 17 00:00:00 2001 From: ThomasX Date: Fri, 8 May 2026 01:41:53 +0800 Subject: [PATCH] fix: treat CP HiCache metadata as backup state --- python/sglang/srt/mem_cache/hiradix_cache.py | 6 +-- .../mem_cache/test_cp_hicache_metadata.py | 45 +++++++++++++++++++ 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index d054acf0e..e7c4f1819 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1006,7 +1006,7 @@ class HiRadixCache(RadixCache): if self._is_pinned(x): # Still active: demote to host if possible - if x.backuped: + if self._node_backuped(x): num_evicted += self._evict_backuped(x) continue written = self.write_backup(x, write_back=True) @@ -1024,7 +1024,7 @@ class HiRadixCache(RadixCache): # Expired pin: clear and fall through to normal eviction self._clear_pin(x) - if not x.backuped: + if not self._node_backuped(x): if self.cache_controller.write_policy == "write_back": # write to host if the node is not backuped num_evicted += self.write_backup(x, write_back=True) @@ -1047,7 +1047,7 @@ class HiRadixCache(RadixCache): if self.cache_controller.write_policy == "write_back": self.writing_check(write_back=True) for node in write_back_nodes: - assert node.backuped + assert self._node_backuped(node) self._evict_backuped(node) self.update_eviction_metrics(num_evicted, start_time) 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 3938955b3..e9ddbad8d 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -10,6 +10,7 @@ for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"): if _mod not in sys.modules: sys.modules[_mod] = MagicMock() +from sglang.srt.mem_cache.base_prefix_cache import EvictParams from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode from sglang.test.ci.ci_register import register_cpu_ci @@ -183,6 +184,13 @@ class FakeEvictionStrategy: return 0 +class FakeEvictDeviceController: + write_policy = "write_through" + + def evict_device(self, device_indices): + return len(device_indices) + + class TestHiRadixCacheCPBackup(CustomTestCase): def test_node_backuped_uses_cp_metadata(self): cache = HiRadixCache.__new__(HiRadixCache) @@ -260,6 +268,43 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertNotIn(1, root.children) self.assertEqual(node.host_len, 16) + def test_evict_demotes_cp_backed_node_without_deleting_radix_child(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.cache_controller = FakeEvictDeviceController() + cache.evictable_leaves = set() + cache.evictable_host_leaves = set() + cache.eviction_strategy = FakeEvictionStrategy() + cache.evictable_size_ = 4 + cache.protected_size_ = 0 + cache._record_remove_event = lambda node: (_ for _ in ()).throw( + AssertionError("must not delete backed radix child") + ) + + root = TreeNode() + root.key = RadixKey(token_ids=[], extra_key=None) + root.value = [] + cache.root_node = root + node = TreeNode() + node.parent = root + node.key = RadixKey(token_ids=[1, 2, 3, 4], extra_key=None) + node.value = torch.arange(4, dtype=torch.int64) + node.host_len = 4 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=4, + owned_positions=torch.tensor([0], dtype=torch.int64), + host_indices=torch.tensor([55], dtype=torch.int64), + ) + root.children[1] = node + cache.evictable_leaves.add(node) + + cache.evict(EvictParams(num_tokens=4)) + + self.assertIn(1, root.children) + self.assertIsNone(node.value) + self.assertIsNotNone(node.cp_hicache) + self.assertEqual(node.cp_hicache.host_indices.tolist(), [55]) + if __name__ == "__main__": unittest.main()