diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 1691fe3b9..dd6ddc290 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -399,10 +399,9 @@ class HiRadixCache(RadixCache): num_evicted += self.cache_controller.evict_host(x.host_value) - for k, v in x.parent.children.items(): - if v == x: - break - del x.parent.children[k] + 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 len(x.parent.children) == 0 and x.parent.evicted: new_priority = self.eviction_strategy.get_priority(x.parent) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 239f37bdc..8ab4ac9a8 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -918,10 +918,10 @@ class MambaRadixCache(BasePrefixCache): node.mamba_value is not None ), f"Invariant violated: leaf node is a tombstone, {node.id=}" assert len(node.children) == 0, f"leaf node has children, {node.id=}" - for k, v in node.parent.children.items(): - if v == node: - break - del node.parent.children[k] + key = self.get_child_key_fn(node.key) + v = node.parent.children.pop(key, None) + assert v == node, f"parent does not have child key, {key}" + self.full_evictable_size_ -= len(node.key) self.mamba_evictable_size_ -= len(node.mamba_value) @@ -935,10 +935,10 @@ class MambaRadixCache(BasePrefixCache): node.mamba_value is None ), f"Deleting a unexpected non-tombstone leaf node, {node.id=}" assert len(node.children) == 0, f"leaf node has children, {node.id=}" - for k, v in node.parent.children.items(): - if v == node: - break - del node.parent.children[k] + key = self.get_child_key_fn(node.key) + v = node.parent.children.pop(key, None) + assert v == node, f"parent does not have child key, {key}" + self.full_evictable_size_ -= len(node.key) def _collect_leaves(self) -> List[TreeNode]: diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 20006b6a9..de6581846 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -730,10 +730,10 @@ class RadixCache(BasePrefixCache): ), f"{key=}, {self.get_child_key_fn(child.key)=}" def _delete_leaf(self, node): - for k, v in node.parent.children.items(): - if v == node: - break - del node.parent.children[k] + key = self.get_child_key_fn(node.key) + v = node.parent.children.pop(key, None) + assert v == node, f"parent does not have child key, {key}" + self.evictable_size_ -= len(node.key) def _total_size_helper(self): diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 7963e4583..3cbc85c7f 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -998,10 +998,9 @@ class SWARadixCache(BasePrefixCache): not node.swa_tombstone ), f"Invariant violated: leaf node is a tombstone, {node.id=}" assert len(node.children) == 0, f"leaf node has children, {node.id=}" - for k, v in node.parent.children.items(): - if v == node: - break - del node.parent.children[k] + key = self.get_child_key_fn(node.key) + v = node.parent.children.pop(key, None) + assert v == node, f"parent does not have child key, {key}" self.full_evictable_size_ -= len(node.key) self.swa_evictable_size_ -= len(node.key) @@ -1015,10 +1014,10 @@ class SWARadixCache(BasePrefixCache): node.swa_tombstone ), f"Deleting a unexpected non-tombstone leaf node, {node.id=}" assert len(node.children) == 0, f"leaf node has children, {node.id=}" - for k, v in node.parent.children.items(): - if v == node: - break - del node.parent.children[k] + key = self.get_child_key_fn(node.key) + v = node.parent.children.pop(key, None) + assert v == node, f"parent does not have child key, {key}" + self.full_evictable_size_ -= len(node.key) def _collect_leaves(self) -> List[TreeNode]: