From 2df832dec0e5ba0534be06646f1683bb5c882ab1 Mon Sep 17 00:00:00 2001 From: ThomasX Date: Fri, 8 May 2026 01:59:14 +0800 Subject: [PATCH] fix: unblock CP host eviction past stale leaves --- python/sglang/srt/mem_cache/hiradix_cache.py | 5 ++ .../mem_cache/test_cp_hicache_metadata.py | 52 +++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 0809b49ef..505e0afd3 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1102,6 +1102,11 @@ class HiRadixCache(RadixCache): continue if not self._node_backuped(x): + if len(x.children) == 0: + parent = self._remove_host_leaf(x) + if parent.evicted and self._node_backuped(parent): + new_priority = self.eviction_strategy.get_priority(parent) + heapq.heappush(eviction_heap, (new_priority, parent)) continue if x.pin_expiry > 0 and time.monotonic() > x.pin_expiry: 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 5aabf1599..4ab7f0827 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -379,6 +379,58 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): self.assertEqual(node.host_len, 0) self.assertIsNone(node.cp_hicache) + def test_cp_host_eviction_unlinks_stale_leaf_to_free_parent_slots(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 + freed = [] + cache.cache_controller = type( + "Controller", + (), + { + "evict_host": lambda self, indices: freed.append(indices.clone()) + or len(indices) + }, + )() + + parent = TreeNode() + parent.parent = cache.root_node + parent.key = RadixKey([1]) + parent.value = None + parent.host_len = 4 + parent.cp_hicache = CpHiCacheNodeMetadata( + logical_len=4, + owned_positions=torch.tensor([2], dtype=torch.int64), + host_indices=torch.tensor([80], dtype=torch.int64), + ) + cache.root_node.children[1] = parent + + stale_child = TreeNode() + stale_child.parent = parent + stale_child.key = RadixKey([2]) + stale_child.value = None + stale_child.host_len = 0 + stale_child.cp_hicache = None + parent.children[2] = stale_child + cache.evictable_host_leaves.add(stale_child) + + physical_freed = cache._evict_host_for_physical_slots(1) + + self.assertEqual(physical_freed, 1) + self.assertEqual(freed[0].tolist(), [80]) + self.assertNotIn(2, parent.children) + self.assertEqual(parent.host_len, 0) + self.assertIsNone(parent.cp_hicache) + self.assertNotIn(1, cache.root_node.children) + if __name__ == "__main__": unittest.main()