From f1a7e605e8fe0d25b658a4973cc769f9cefe75a1 Mon Sep 17 00:00:00 2001 From: ThomasX Date: Fri, 8 May 2026 01:50:07 +0800 Subject: [PATCH] feat: split and evict CP HiCache metadata --- python/sglang/srt/mem_cache/hiradix_cache.py | 18 ++++- .../mem_cache/test_cp_hicache_metadata.py | 74 +++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index e7c4f1819..0809b49ef 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1084,6 +1084,9 @@ class HiRadixCache(RadixCache): return parent def _evict_host_for_physical_slots(self, required_host_slots: int) -> int: + if required_host_slots <= 0: + return 0 + leaves = list(self.evictable_host_leaves) eviction_heap = [ (self.eviction_strategy.get_priority(node), node) for node in leaves @@ -1098,6 +1101,9 @@ class HiRadixCache(RadixCache): if not x.evicted: continue + if not self._node_backuped(x): + continue + if x.pin_expiry > 0 and time.monotonic() > x.pin_expiry: self._clear_pin(x) @@ -1124,6 +1130,9 @@ class HiRadixCache(RadixCache): return num_evicted def evict_host(self, num_tokens: int): + if self._uses_cp_hicache: + return self._evict_host_for_physical_slots(num_tokens) + leaves = list(self.evictable_host_leaves) eviction_heap = [ (self.eviction_strategy.get_priority(node), node) for node in leaves @@ -1596,7 +1605,14 @@ class HiRadixCache(RadixCache): else: new_node.value = child.value[:split_len].clone() child.value = child.value[split_len:].clone() - if child.backuped: + if self._uses_cp_hicache: + if self._node_backuped(child): + new_node.cp_hicache, child.cp_hicache = child.cp_hicache.split( + split_len + ) + new_node.host_len = split_len + child.host_len = child.host_len - split_len + elif child.backuped: new_node.host_value = child.host_value[:split_len].clone() child.host_value = child.host_value[split_len:].clone() 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 e9ddbad8d..5aabf1599 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -306,5 +306,79 @@ class TestHiRadixCacheCPBackup(CustomTestCase): self.assertEqual(node.cp_hicache.host_indices.tolist(), [55]) +class TestHiRadixCacheCPSplitEvict(CustomTestCase): + def test_split_node_splits_cp_metadata_by_owned_positions(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache.page_size = 1 + + root = TreeNode() + root.key = RadixKey([]) + child = TreeNode() + child.parent = root + child.key = RadixKey(list(range(10))) + child.value = None + child.hash_value = None + child.host_len = 10 + child.cp_hicache = CpHiCacheNodeMetadata( + logical_len=10, + owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64), + host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64), + ) + root.children[0] = child + + new_node = cache._split_node(child.key, child, 5) + + self.assertEqual(new_node.host_len, 5) + self.assertEqual(child.host_len, 5) + self.assertEqual(new_node.cp_hicache.owned_positions.tolist(), [0, 2]) + self.assertEqual(new_node.cp_hicache.host_indices.tolist(), [20, 21]) + self.assertEqual(child.cp_hicache.owned_positions.tolist(), [0, 4]) + self.assertEqual(child.cp_hicache.host_indices.tolist(), [22, 23]) + + def test_cp_host_eviction_uses_physical_freed_slots_for_progress(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + 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 + cache._update_host_leaf_status = 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.tensor([1], dtype=torch.int64), + host_indices=torch.tensor([70], dtype=torch.int64), + ) + cache.root_node.children[1] = node + cache.evictable_host_leaves.add(node) + + physical_freed = cache._evict_host_for_physical_slots(1) + + self.assertEqual(physical_freed, 1) + self.assertEqual(freed[0].tolist(), [70]) + self.assertEqual(node.host_len, 0) + self.assertIsNone(node.cp_hicache) + + if __name__ == "__main__": unittest.main()