fix merge efficiency bug

Author: shihaobai

.pre-commit-config.yaml

@@ -7,8 +7,7 @@ repos:
         args: [--line-length=120]
         additional_dependencies: ['click==8.0.4']
 -   repo: https://github.com/pycqa/flake8
-    rev: 3.9.0
+    rev: 6.1.0 
     hooks:
       - id: flake8
-        additional_dependencies: [flake8-typing-imports==1.9.0]
-        args: ['--config=.flake8', '--max-line-length=120', '--ignore=TYP001, E722, C901, E203, E266, E402, E302, E241, E902, E731, F403, E701, F405, F401, W292, W293, W503, W606, E231']
+        args: ['--max-line-length=120', '--ignore=TYP001, E722, C901, E203, E266, E402, E302, E241, E902, E731, F403, E701, F405, F401, W292, W293, W503, W606, E231']

lightllm/server/router/dynamic_prompt/radix_cache.py

@@ -342,7 +342,7 @@ def evict(self, need_remove_tokens, evict_callback):
 
         return
 
-    def _try_merge(self, child_node: TreeNode):
+    def _try_merge(self, child_node: TreeNode) -> Optional[TreeNode]:
         """
         合并条件:
         1. 父节点不是根节点。
@@ -352,52 +352,54 @@ def _try_merge(self, child_node: TreeNode):
         """
         parent_node = child_node.parent
         # 条件检查
-        if parent_node is None or parent_node == self.root_node:
-            return False
-        
-        if parent_node.ref_counter == 0 and \
-           child_node.ref_counter == 0 and \
-           len(parent_node.children) == 1:
-            
-            if child_node.is_leaf():
-                self.evict_tree_set.discard(child_node)
-
-            new_key = torch.cat([parent_node.token_id_key, child_node.token_id_key])
-            new_value = torch.cat([parent_node.token_mem_index_value, child_node.token_mem_index_value])
-
-            parent_node.token_id_key = new_key
-            parent_node.token_mem_index_value = new_value
-            parent_node.children = child_node.children
-            for grandchild in parent_node.children.values():
-                grandchild.parent = parent_node
-
-            parent_node.node_value_len = len(parent_node.token_mem_index_value)
-            parent_node.time_id = max(parent_node.time_id, child_node.time_id)
+        if (
+            parent_node is None
+            or parent_node == self.root_node
+            or parent_node.ref_counter != 0
+            or len(parent_node.children) != 1
+            or child_node.ref_counter != 0
+        ):
+            return None
 
-            if parent_node.is_leaf():
-                self.evict_tree_set.add(parent_node)
-            
-            child_node.parent = None
-            return True
+        if child_node.is_leaf():
+            self.evict_tree_set.discard(child_node)
+
+        child_node.token_id_key = torch.cat([parent_node.token_id_key, child_node.token_id_key])
+        child_node.token_mem_index_value = torch.cat(
+            [parent_node.token_mem_index_value, child_node.token_mem_index_value]
+        )
+        child_node.node_value_len = len(child_node.token_mem_index_value)
+        child_node.time_id = max(parent_node.time_id, child_node.time_id)
+
+        grandparent_node = parent_node.parent
+        key_in_grandparent = parent_node.token_id_key[0].item()
+        grandparent_node.children[key_in_grandparent] = child_node
+        child_node.parent = grandparent_node
 
-        return False
+        parent_node.parent = None
+
+        if child_node.is_leaf():
+            self.evict_tree_set.add(child_node)
+
+        return child_node
 
     def merge_unreferenced_nodes(self):
-        if not self.root_node.children:
-            return
-        nodes_to_process = []
-        traversal_stack = list(self.root_node.children.values())
-
-        while traversal_stack:
-            node = traversal_stack.pop()
-            nodes_to_process.append(node)
-            traversal_stack.extend(list(node.children.values()))
-
-        nodes_to_process.reverse()
-        for node in nodes_to_process:
-            if node.parent is not None:
-                self._try_merge(node)
-    
+        worklist = collections.deque(
+            [
+                node
+                for node in self.evict_tree_set
+                if node.ref_counter == 0 and node.parent is not None and node.parent != self.root_node
+            ]
+        )
+
+        while worklist:
+            node = worklist.popleft()
+            if node.parent is None:
+                continue
+            merged_node = self._try_merge(node)
+            if merged_node:
+                worklist.append(merged_node)
+
     def assert_leafs_is_right(self):
         for node in self.evict_tree_set:
             if node.is_leaf() and node.ref_counter == 0:

unit_tests/server/router/dynamic_prompt/test_radix_cache.py

@@ -97,7 +97,7 @@ def test_case5():
     """
     print("\nTest Case 5: Merging simple parent-child nodes when ref_counter is 0\n")
     tree = RadixCache("unique_name", 100, 0)
-    
+
     _, node_a = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64))
     _, node_b = tree.insert(torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64))
     tree.print_self()
@@ -114,10 +114,11 @@ def test_case5():
     tree.merge_unreferenced_nodes()
     tree.print_self()
 
-    assert torch.equal(node_a.token_id_key, torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64))
-    assert node_a.is_leaf()
+    assert torch.equal(node_b.token_id_key, torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64))
+    assert node_b.is_leaf()
     assert tree.get_tree_total_tokens_num() == 5
-    assert tree.root_node.children[1] is node_a
+    assert tree.root_node.children[1] is node_b
+
 
 def test_case6():
     """
@@ -137,18 +138,19 @@ def test_case6():
     tree.print_self()
 
     assert len(tree.root_node.children) == 1
-    # 节点 A 的 key 应该是完整的 [1, 2, 3, 4]
-    assert torch.equal(node_a.token_id_key, torch.tensor([1, 2, 3, 4], dtype=torch.int64))
-    assert node_a.is_leaf()
+    # 节点 C 的 key 应该是完整的 [1, 2, 3, 4]
+    assert torch.equal(node_c.token_id_key, torch.tensor([1, 2, 3, 4], dtype=torch.int64))
+    assert node_c.is_leaf()
     assert tree.get_tree_total_tokens_num() == 4
 
+
 def test_case7():
     """
     测试场景：由于父节点或子节点的 ref_counter > 0，合并不应该发生。
     """
     print("\nTest Case 7: Merging when parent or child ref_counter > 0\n")
     tree = RadixCache("unique_name", 100, 0)
-    
+
     _, node_a = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64))
     _, node_b = tree.insert(torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64))
     tree.print_self()
@@ -157,14 +159,15 @@ def test_case7():
     assert matched_node is node_a
     assert node_a.ref_counter == 1
     assert node_b.ref_counter == 0
-    
+
     tree.merge_unreferenced_nodes()
     tree.print_self()
 
     assert torch.equal(node_a.token_id_key, torch.tensor([1, 2, 3], dtype=torch.int64))
     assert not node_a.is_leaf()
     assert node_b.parent is node_a
 
+
 def test_case8():
     """
     测试场景：由于父节点有多个子节点，合并不应该发生。
@@ -184,12 +187,13 @@ def test_case8():
 
     tree.merge_unreferenced_nodes()
     tree.print_self()
-    
+
     assert len(node_a.children) == 2
     assert torch.equal(node_a.token_id_key, torch.tensor([1, 2], dtype=torch.int64))
     assert tree.root_node.children[1].children[3] is node_b
     assert tree.root_node.children[1].children[4] is node_c
 
+
 def test_case9():
     """
     测试场景：在一个复杂的树中，只有满足条件的分支被合并。
@@ -204,7 +208,7 @@ def test_case9():
     # 分支2: 不可合并的链 C -> D (因为 C 被引用)
     _, node_c = tree.insert(torch.tensor([4, 5], dtype=torch.int64))
     _, node_d = tree.insert(torch.tensor([4, 5, 6], dtype=torch.int64))
-    
+
     # 增加 C 的引用计数
     tree.match_prefix(torch.tensor([4, 5], dtype=torch.int64), update_refs=True)
     assert node_c.ref_counter == 1
@@ -213,15 +217,15 @@ def test_case9():
     tree.merge_unreferenced_nodes()
     tree.print_self()
 
-    merged_node_a = tree.root_node.children[1]
-    assert torch.equal(merged_node_a.token_id_key, torch.tensor([1, 2, 3], dtype=torch.int64))
-    assert merged_node_a.is_leaf()
-    
+    merged_node_b = tree.root_node.children[1]
+    assert torch.equal(merged_node_b.token_id_key, torch.tensor([1, 2, 3], dtype=torch.int64))
+    assert merged_node_b.is_leaf()
+
     unmerged_node_c = tree.root_node.children[4]
     assert torch.equal(unmerged_node_c.token_id_key, torch.tensor([4, 5], dtype=torch.int64))
     assert not unmerged_node_c.is_leaf()
     assert len(unmerged_node_c.children) == 1
-    
+
     unmerged_node_d = unmerged_node_c.children[6]
     assert torch.equal(unmerged_node_d.token_id_key, torch.tensor([6], dtype=torch.int64))