fix radix cache

Author: wangzaijun

lightllm/distributed/custom_all_gather.py

@@ -28,7 +28,6 @@
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.device_utils import has_nvlink
 from lightllm.utils.light_utils import light_ops
-from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
 
 
 try:

lightllm/distributed/custom_all_reduce.py

@@ -29,7 +29,6 @@
 from lightllm.utils.device_utils import has_nvlink
 from lightllm.utils.sgl_utils import sgl_allreduce_ops
 from lightllm.utils.vllm_utils import vllm_ops
-from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
 
 logger = init_logger(__name__)
 
@@ -225,6 +224,9 @@ def all_reduce(self, inp: torch.Tensor, *, out: torch.Tensor = None, registered:
         buffer.
         """
         if out is None:
+            # fix circle import
+            from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
+
             out = g_cache_manager.alloc_tensor(inp.shape, inp.dtype, device=inp.device, is_graph_out=False)
         if registered:
             ops.all_reduce(self._ptr, inp, out, 0, 0)
@@ -243,6 +245,9 @@ def custom_all_reduce(self, input: torch.Tensor) -> Optional[torch.Tensor]:
             else:
                 # If warm up, mimic the allocation pattern since custom
                 # allreduce is out-of-place.
+                # fix circle import
+                from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
+
                 out = g_cache_manager.alloc_tensor(input.shape, input.dtype, device=input.device, is_graph_out=False)
                 return out
         else:

lightllm/server/router/dynamic_prompt/radix_cache.py

@@ -1,10 +1,10 @@
 # Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/managers/router/radix_cache.py
 import torch
 import numpy as np
-from typing import Tuple, Dict, Set, List, Optional
+import collections
+from typing import Tuple, Dict, Set, List, Optional, Union
 from sortedcontainers import SortedSet
 from .shared_arr import SharedArray
-from lightllm.common.mem_manager import MemoryManager
 
 
 class UniqueTimeIdGenerator:
@@ -103,8 +103,10 @@ class RadixCache:
     unique_name 主要用于解决单机，多实列部署时的shm冲突
     """
 
-    def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager: MemoryManager = None):
-        self.mem_manager = mem_manager
+    def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager=None):
+        from lightllm.common.mem_manager import MemoryManager
+
+        self.mem_manager: MemoryManager = mem_manager
         self._key_dtype = torch.int64
         self._value_dtype = torch.int64
 
@@ -133,58 +135,100 @@ def insert(self, key, value=None) -> Tuple[int, Optional[TreeNode]]:
         return self._insert_helper(self.root_node, key, value)
 
     def _insert_helper(self, node: TreeNode, key, value) -> Tuple[int, Optional[TreeNode]]:
+        handle_stack = collections.deque()
+        update_list = collections.deque()
+        handle_stack.append((node, key, value))
+
+        ans_prefix_len = 0
+        ans_node = None
+
+        while len(handle_stack) != 0:
+            node, key, value = handle_stack.popleft()
+            ans_tuple = self._insert_helper_no_recursion(node=node, key=key, value=value)
+            if len(ans_tuple) == 4:
+                (_prefix_len, new_node, new_key, new_value) = ans_tuple
+                ans_prefix_len += _prefix_len
+                handle_stack.append((new_node, new_key, new_value))
+            else:
+                _prefix_len, ans_node = ans_tuple
+                ans_prefix_len += _prefix_len
+
+            update_list.append(node)
+
+        while len(update_list) != 0:
+            cur_node: TreeNode = update_list.pop()
+            cur_node.update_time()
+            if cur_node.is_leaf():
+                self.evict_tree_set.add(cur_node)
+
+        assert ans_node is not None
+
+        return ans_prefix_len, ans_node
+
+    def _insert_helper_no_recursion(
+        self, node: TreeNode, key: torch.Tensor, value: torch.Tensor
+    ) -> Union[Tuple[int, Optional[TreeNode]], Tuple[int, TreeNode, torch.Tensor, torch.Tensor]]:
         if node.is_leaf():
             self.evict_tree_set.discard(node)
 
-        try:
-            first_key_id = key[0].item()
-            if first_key_id in node.children.keys():
-                child: TreeNode = node.children[first_key_id]
-                prefix_len = match(key, child.token_id_key)
-                if prefix_len == len(key):
+        first_key_id = key[0].item()
+        if first_key_id in node.children.keys():
+            child: TreeNode = node.children[first_key_id]
+            prefix_len = match(key, child.token_id_key)
+            if prefix_len == len(key):
+                if prefix_len == len(child.token_id_key):
                     if child.is_leaf():
                         self.evict_tree_set.discard(child)
                     child.update_time()
                     if child.is_leaf():
                         self.evict_tree_set.add(child)
                     return prefix_len, child
-
-                elif prefix_len < len(key) and prefix_len < len(child.token_id_key):
+                elif prefix_len < len(child.token_id_key):
                     if child.is_leaf():
                         self.evict_tree_set.discard(child)
 
-                    key = key[prefix_len:]
-                    value = value[prefix_len:]
                     split_parent_node = child.split_node(prefix_len)
-                    new_node = split_parent_node.add_and_return_new_child(key, value)
-                    # update total token num
-                    self.tree_total_tokens_num.arr[0] += len(new_node.token_mem_index_value)
 
                     if split_parent_node.is_leaf():
                         self.evict_tree_set.add(split_parent_node)
-                    if new_node.is_leaf():
-                        self.evict_tree_set.add(new_node)
-
                     if child.is_leaf():
                         self.evict_tree_set.add(child)
-                    return prefix_len, new_node
-                elif prefix_len < len(key) and prefix_len == len(child.token_id_key):
-                    _prefix_len, ans_node = self._insert_helper(child, key[prefix_len:], value[prefix_len:])
-                    return prefix_len + _prefix_len, ans_node
+
+                    return prefix_len, split_parent_node
                 else:
                     assert False, "can not run to here"
 
-            else:
-                new_node = node.add_and_return_new_child(key, value)
+            elif prefix_len < len(key) and prefix_len < len(child.token_id_key):
+                if child.is_leaf():
+                    self.evict_tree_set.discard(child)
+
+                key = key[prefix_len:]
+                value = value[prefix_len:]
+                split_parent_node = child.split_node(prefix_len)
+                new_node = split_parent_node.add_and_return_new_child(key, value)
                 # update total token num
                 self.tree_total_tokens_num.arr[0] += len(new_node.token_mem_index_value)
+
+                if split_parent_node.is_leaf():
+                    self.evict_tree_set.add(split_parent_node)
                 if new_node.is_leaf():
                     self.evict_tree_set.add(new_node)
-                return 0, new_node
-        finally:
-            node.update_time()
-            if node.is_leaf():
-                self.evict_tree_set.add(node)
+
+                if child.is_leaf():
+                    self.evict_tree_set.add(child)
+                return prefix_len, new_node
+            elif prefix_len < len(key) and prefix_len == len(child.token_id_key):
+                return (prefix_len, child, key[prefix_len:], value[prefix_len:])
+            else:
+                assert False, "can not run to here"
+
+        else:
+            new_node = node.add_and_return_new_child(key, value)
+            # update total token num
+            self.tree_total_tokens_num.arr[0] += len(new_node.token_mem_index_value)
+            if new_node.is_leaf():
+                self.evict_tree_set.add(new_node)
+            return 0, new_node
 
     def match_prefix(self, key, update_refs=False):
         assert len(key) != 0
@@ -200,7 +244,39 @@ def match_prefix(self, key, update_refs=False):
             self.dec_node_ref_counter(self.root_node)
             return None, 0, None
 
-    def _match_prefix_helper(self, node: TreeNode, key, ans_value_list: list, update_refs=False) -> TreeNode:
+    def _match_prefix_helper(
+        self, node: TreeNode, key: torch.Tensor, ans_value_list: list, update_refs=False
+    ) -> TreeNode:
+        handle_stack = collections.deque()
+        update_list = collections.deque()
+        handle_stack.append((node, key))
+
+        ans_node = None
+
+        while len(handle_stack) != 0:
+            node, key = handle_stack.popleft()
+            ans_tuple = self._match_prefix_helper_no_recursion(
+                node=node, key=key, ans_value_list=ans_value_list, update_refs=update_refs
+            )
+            if isinstance(ans_tuple, tuple):
+                new_node, new_key = ans_tuple
+                handle_stack.append((new_node, new_key))
+            else:
+                ans_node = ans_tuple
+
+            update_list.append(node)
+
+        while len(update_list) != 0:
+            cur_node: TreeNode = update_list.pop()
+            cur_node.update_time()
+            if cur_node.is_leaf():
+                self.evict_tree_set.add(cur_node)
+
+        return ans_node
+
+    def _match_prefix_helper_no_recursion(
+        self, node: TreeNode, key: torch.Tensor, ans_value_list: list, update_refs=False
+    ) -> TreeNode:
         if node.is_leaf():
             self.evict_tree_set.discard(node)
 
@@ -210,44 +286,39 @@ def _match_prefix_helper(self, node: TreeNode, key, ans_value_list: list, update
             if node.ref_counter == 1:
                 self.refed_tokens_num.arr[0] += len(node.token_mem_index_value)
 
-        try:
-            if len(key) == 0:
-                return node
+        if len(key) == 0:
+            return node
 
-            first_key_id = key[0].item()
-            if first_key_id not in node.children.keys():
-                return node
+        first_key_id = key[0].item()
+        if first_key_id not in node.children.keys():
+            return node
+        else:
+            child = node.children[first_key_id]
+            prefix_len = match(key, child.token_id_key)
+            if prefix_len == len(child.token_id_key):
+                ans_value_list.append(child.token_mem_index_value)
+                return (child, key[prefix_len:])
+            elif prefix_len < len(child.token_id_key):
+                if child.is_leaf():
+                    self.evict_tree_set.discard(child)
+
+                split_parent_node = child.split_node(prefix_len)
+                ans_value_list.append(split_parent_node.token_mem_index_value)
+
+                if update_refs:
+                    split_parent_node.ref_counter += 1
+                    # from 0 to 1 need update refs token num
+                    if split_parent_node.ref_counter == 1:
+                        self.refed_tokens_num.arr[0] += len(split_parent_node.token_mem_index_value)
+
+                if child.is_leaf():
+                    self.evict_tree_set.add(child)
+                if split_parent_node.is_leaf():
+                    self.evict_tree_set.add(split_parent_node)
+
+                return split_parent_node
             else:
-                child = node.children[first_key_id]
-                prefix_len = match(key, child.token_id_key)
-                if prefix_len == len(child.token_id_key):
-                    ans_value_list.append(child.token_mem_index_value)
-                    return self._match_prefix_helper(child, key[prefix_len:], ans_value_list, update_refs=update_refs)
-                elif prefix_len < len(child.token_id_key):
-                    if child.is_leaf():
-                        self.evict_tree_set.discard(child)
-
-                    split_parent_node = child.split_node(prefix_len)
-                    ans_value_list.append(split_parent_node.token_mem_index_value)
-
-                    if update_refs:
-                        split_parent_node.ref_counter += 1
-                        # from 0 to 1 need update refs token num
-                        if split_parent_node.ref_counter == 1:
-                            self.refed_tokens_num.arr[0] += len(split_parent_node.token_mem_index_value)
-
-                    if child.is_leaf():
-                        self.evict_tree_set.add(child)
-                    if split_parent_node.is_leaf():
-                        self.evict_tree_set.add(split_parent_node)
-
-                    return split_parent_node
-                else:
-                    assert False, "error state"
-        finally:
-            node.update_time()
-            if node.is_leaf():
-                self.evict_tree_set.add(node)
+                assert False, "error state"
 
     def evict(self, need_remove_tokens, evict_callback):
         if self.tree_total_tokens_num.arr[0] - self.refed_tokens_num.arr[0] < need_remove_tokens:
@@ -417,3 +488,7 @@ def get_tree_total_tokens_num(self, dp_rank_in_node):
 
     def get_unrefed_tokens_num(self, dp_rank_in_node):
         return self.dp_rank_clients[dp_rank_in_node].get_unrefed_tokens_num()
+
+
+class _RecursionParams:
+    pass

lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py

@@ -165,6 +165,7 @@ def _master_req_to_radix_cache(self, master_req: InferReq):
         value = self.model.req_manager.req_to_token_indexs[master_req.req_idx][: master_req.cur_kv_len].detach().cpu()
         prefix_len, new_shared_kv_node = self.radix_cache.insert(key, value)
         old_prefix_len = 0 if master_req.shared_kv_node is None else master_req.shared_kv_node.node_prefix_total_len
+        assert old_prefix_len <= master_req.cur_kv_len
         self.model.mem_manager.free(
             self.model.req_manager.req_to_token_indexs[master_req.req_idx][old_prefix_len:prefix_len]
         )
@@ -173,7 +174,9 @@ def _master_req_to_radix_cache(self, master_req: InferReq):
         self.radix_cache.dec_node_ref_counter(master_req.shared_kv_node)
         self.radix_cache.add_node_ref_counter(new_shared_kv_node)
         master_req.shared_kv_node = new_shared_kv_node
-        assert new_shared_kv_node.node_prefix_total_len == master_req.cur_kv_len
+        assert (
+            new_shared_kv_node.node_prefix_total_len == master_req.cur_kv_len
+        ), f"shared len: {new_shared_kv_node.node_prefix_total_len} cur_kv_len {master_req.cur_kv_len}"
 
         share_node, kv_len, value = self.radix_cache.match_prefix(key, update_refs=False)
         assert share_node == new_shared_kv_node and kv_len == master_req.cur_kv_len