impl radix_cache node merge

Author: Longxmas

lightllm/server/router/dynamic_prompt/radix_cache.py

@@ -342,6 +342,62 @@ def evict(self, need_remove_tokens, evict_callback):
 
         return
 
+    def _try_merge(self, child_node: TreeNode):
+        """
+        合并条件:
+        1. 父节点不是根节点。
+        2. 父节点的引用计数为 0。
+        3. 子节点的引用计数为 0。
+        4. 父节点只有一个子节点 (即 child_node)。
+        """
+        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_leaf():
+                self.evict_tree_set.add(parent_node)
+            
+            child_node.parent = None
+            return True
+
+        return False
+
+    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)
+    
     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

@@ -91,5 +91,140 @@ def test_case4():
     return
 
 
+def test_case5():
+    """
+    测试场景：一个简单的父子节点链 (A -> B)，在 ref_counter 都为 0 时，应该成功合并。
+    """
+    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()
+
+    # 验证初始状态：A -> B 结构，且 ref_counter 均为 0
+    assert node_b.parent == node_a
+    assert torch.equal(node_a.token_id_key, torch.tensor([1, 2, 3], dtype=torch.int64))
+    assert len(node_a.children) == 1
+    assert node_a.ref_counter == 0
+    assert node_b.ref_counter == 0
+    assert tree.get_tree_total_tokens_num() == 5
+
+    # 执行合并
+    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 tree.get_tree_total_tokens_num() == 5
+    assert tree.root_node.children[1] is node_a
+
+def test_case6():
+    """
+    测试场景：一个长的节点链 (A -> B -> C)，在 ref_counter 都为 0 时，应该级联合并成一个节点。
+    """
+    print("\nTest Case 6: Merging long nodes when ref_counter is 0\n")
+    tree = RadixCache("unique_name", 100, 0)
+    _, node_a = tree.insert(torch.tensor([1], dtype=torch.int64))
+    _, node_b = tree.insert(torch.tensor([1, 2], dtype=torch.int64))
+    _, node_c = tree.insert(torch.tensor([1, 2, 3, 4], dtype=torch.int64))
+    tree.print_self()
+
+    assert node_c.parent == node_b
+    assert node_b.parent == node_a
+    assert tree.get_tree_total_tokens_num() == 4
+    tree.merge_unreferenced_nodes()
+    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()
+    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()
+
+    matched_node, _, _ = tree.match_prefix(torch.tensor([1, 2, 3], dtype=torch.int64), update_refs=True)
+    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():
+    """
+    测试场景：由于父节点有多个子节点，合并不应该发生。
+    """
+    print("\nTest Case 8: Merging when parent has multiple children\n")
+    tree = RadixCache("unique_name", 100, 0)
+
+    _, node_a = tree.insert(torch.tensor([1, 2], dtype=torch.int64))
+    _, node_b = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64))
+    _, node_c = tree.insert(torch.tensor([1, 2, 4], dtype=torch.int64))
+    tree.print_self()
+
+    assert len(node_a.children) == 2
+    assert node_a.ref_counter == 0
+    assert node_b.ref_counter == 0
+    assert node_c.ref_counter == 0
+
+    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():
+    """
+    测试场景：在一个复杂的树中，只有满足条件的分支被合并。
+    """
+    print("\nTest Case 9: Merging in a complex tree with mixed conditions\n")
+    tree = RadixCache("unique_name", 100, 0)
+
+    # 分支1: 可合并的链 A -> B
+    _, node_a = tree.insert(torch.tensor([1, 2], dtype=torch.int64))
+    _, node_b = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64))
+
+    # 分支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
+    tree.print_self()
+
+    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()
+    
+    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))
+
+
 if __name__ == "__main__":
     pytest.main()