[RadixTree] Optimize the Time Complexity of Node Retrieval Operation from O(n*m) to O(n) (#13334)

Signed-off-by: CLFutureX <chenyongqyl@163.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>

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():
+            key = self.get_child_key_fn(x.key)
-                if v == x:
+            v = x.parent.children.pop(key, None)
-                    break
+            assert v == x, f"parent does not have child key, {key}"
-            del x.parent.children[k]
 
             if len(x.parent.children) == 0 and x.parent.evicted:
                 new_priority = self.eviction_strategy.get_priority(x.parent)

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():
+        key = self.get_child_key_fn(node.key)
-            if v == node:
+        v = node.parent.children.pop(key, None)
-                break
+        assert v == node, f"parent does not have child key, {key}"
-        del node.parent.children[k]
+
         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():
+        key = self.get_child_key_fn(node.key)
-            if v == node:
+        v = node.parent.children.pop(key, None)
-                break
+        assert v == node, f"parent does not have child key, {key}"
-        del node.parent.children[k]
+
         self.full_evictable_size_ -= len(node.key)
 
     def _collect_leaves(self) -> List[TreeNode]:

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():
+        key = self.get_child_key_fn(node.key)
-            if v == node:
+        v = node.parent.children.pop(key, None)
-                break
+        assert v == node, f"parent does not have child key, {key}"
-        del node.parent.children[k]
+
         self.evictable_size_ -= len(node.key)
 
     def _total_size_helper(self):

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():
+        key = self.get_child_key_fn(node.key)
-            if v == node:
+        v = node.parent.children.pop(key, None)
-                break
+        assert v == node, f"parent does not have child key, {key}"
-        del node.parent.children[k]
         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():
+        key = self.get_child_key_fn(node.key)
-            if v == node:
+        v = node.parent.children.pop(key, None)
-                break
+        assert v == node, f"parent does not have child key, {key}"
-        del node.parent.children[k]
+
         self.full_evictable_size_ -= len(node.key)
 
     def _collect_leaves(self) -> List[TreeNode]: