diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 94a722384..d054acf0e 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1073,6 +1073,56 @@ class HiRadixCache(RadixCache): self._delete_leaf(node) return num_evicted + def _remove_host_leaf(self, node: TreeNode) -> TreeNode: + parent = node.parent + key = self.get_child_key_fn(node.key) + v = parent.children.pop(key, None) + assert v == node, f"parent does not have child key, {key}" + if node in self.evictable_host_leaves: + self.evictable_host_leaves.remove(node) + self._update_host_leaf_status(parent) + return parent + + def _evict_host_for_physical_slots(self, required_host_slots: int) -> int: + leaves = list(self.evictable_host_leaves) + eviction_heap = [ + (self.eviction_strategy.get_priority(node), node) for node in leaves + ] + heapq.heapify(eviction_heap) + + num_evicted = 0 + while num_evicted < required_host_slots and len(eviction_heap): + _priority, x = heapq.heappop(eviction_heap) + if x == self.root_node: + break + if not x.evicted: + continue + + if x.pin_expiry > 0 and time.monotonic() > x.pin_expiry: + self._clear_pin(x) + + if x.host_ref_counter > 0: + continue + + host_indices = ( + x.cp_hicache.host_indices if x.cp_hicache is not None else None + ) + physical_count = len(host_indices) if host_indices is not None else 0 + + self._record_remove_event(x) + if physical_count > 0: + num_evicted += self.cache_controller.evict_host(host_indices) + x.host_len = 0 + x.cp_hicache = None + x.host_value = None + + parent = self._remove_host_leaf(x) + if len(parent.children) == 0 and parent.evicted: + new_priority = self.eviction_strategy.get_priority(parent) + heapq.heappush(eviction_heap, (new_priority, parent)) + + return num_evicted + def evict_host(self, num_tokens: int): leaves = list(self.evictable_host_leaves) eviction_heap = [ @@ -1102,16 +1152,11 @@ class HiRadixCache(RadixCache): self._record_remove_event(x) num_evicted += self.cache_controller.evict_host(x.host_value) - key = self.get_child_key_fn(x.key) - v = x.parent.children.pop(key, None) - assert v == x, f"parent does not have child key, {key}" - if x in self.evictable_host_leaves: - self.evictable_host_leaves.remove(x) - self._update_host_leaf_status(x.parent) + parent = self._remove_host_leaf(x) - if len(x.parent.children) == 0 and x.parent.evicted: - new_priority = self.eviction_strategy.get_priority(x.parent) - heapq.heappush(eviction_heap, (new_priority, x.parent)) + if len(parent.children) == 0 and parent.evicted: + new_priority = self.eviction_strategy.get_priority(parent) + heapq.heappush(eviction_heap, (new_priority, parent)) def load_back( self, node: TreeNode, mem_quota: Optional[int] = None 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 13e06cd89..3938955b3 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -11,7 +11,7 @@ for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"): sys.modules[_mod] = MagicMock() from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache -from sglang.srt.mem_cache.radix_cache import TreeNode +from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -159,6 +159,7 @@ class FakeWriteController: self.required_host_slots = required_host_slots self.calls = 0 self.write_policy = "write_through" + self.evicted_host_indices = [] def write(self, device_indices, node_id=-1, priority=None): self.calls += 1 @@ -172,6 +173,15 @@ class FakeWriteController: ) ) + def evict_host(self, host_indices): + self.evicted_host_indices.append(host_indices.clone()) + return len(host_indices) + + +class FakeEvictionStrategy: + def get_priority(self, node): + return 0 + class TestHiRadixCacheCPBackup(CustomTestCase): def test_node_backuped_uses_cp_metadata(self): @@ -212,18 +222,42 @@ class TestHiRadixCacheCPBackup(CustomTestCase): def test_write_backup_retries_by_required_physical_slots(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True - cache.cache_controller = FakeWriteController(required_host_slots=3) - cache._evict_host_for_physical_slots = lambda required: setattr( - cache, "evicted_required", required - ) + cache.cache_controller = FakeWriteController(required_host_slots=1) + cache.evictable_host_leaves = set() + cache.eviction_strategy = FakeEvictionStrategy() + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache._record_remove_event = lambda node: None cache.ongoing_write_through = {} cache.inc_node_lock_ref = lambda node: None + + root = TreeNode() + root.key = RadixKey(token_ids=[], extra_key=None) + root.value = [] + cache.root_node = root + evictable_node = TreeNode() + evictable_node.parent = root + evictable_node.key = RadixKey(token_ids=[1], extra_key=None) + evictable_node.value = None + evictable_node.host_len = 4 + evictable_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] = evictable_node + cache.evictable_host_leaves.add(evictable_node) + node = TreeNode() node.value = torch.arange(16, dtype=torch.int64) cache.write_backup(node) - self.assertEqual(cache.evicted_required, 3) + self.assertEqual( + cache.cache_controller.evicted_host_indices[0].tolist(), [55] + ) + self.assertEqual(evictable_node.host_len, 0) + self.assertIsNone(evictable_node.cp_hicache) + self.assertNotIn(1, root.children) self.assertEqual(node.host_len, 16)