fix

Author: sufubao

lightllm/server/router/dynamic_prompt/hybrid_radix_cache.py

@@ -5,6 +5,9 @@
 
 from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache, TreeNode
 from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
 
 
 class HybridMemManager(MemoryManager):
@@ -30,6 +33,10 @@ def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager=None)
         super().__init__(unique_name, total_token_num, rank_in_node, mem_manager)
         # 用于缓存需要被驱逐的buffer节点， 应该包含所有有buffer的节点
         self.evict_buffer_set: Set[TreeNode] = SortedSet(key=lambda x: (x.buffer_time,))
+        self.match_count = 0
+        self.log_interval = 1000
+        self.match_len = 0
+        self.hit_len = 0
 
     def free_radix_cache_to_get_enough_buffer(self, need_buffer_num):
         if need_buffer_num > self.mem_manager.get_buffer_can_use_size():
@@ -47,14 +54,14 @@ def release_buffer(buffer_idx):
                 release_buffers.append(buffer_idx)
                 return
 
-            self.evict_buffer(need_evict_buffer_num, release_buffer, release_mem)
+            self._evict_buffer(need_evict_buffer_num, release_buffer, release_mem)
             self.mem_manager.free_buffer(release_buffers)
             if len(release_mems) > 0:
                 mem_index = torch.concat(release_mems)
                 self.mem_manager.free(mem_index)
         return
 
-    def evict_buffer(self, need_evict_buffer_num, evict_buffer_callback, evict_token_callback):
+    def _evict_buffer(self, need_evict_buffer_num, evict_buffer_callback, evict_token_callback):
         while need_evict_buffer_num > 0:
             node = self.evict_buffer_set.pop(0)
             assert node.buffer_idx is not None
@@ -78,6 +85,7 @@ def insert_for_hybrid_radix_cache(self, reqs):
 
         self.free_radix_cache_to_get_enough_buffer(len(reqs))
         new_buffer_indexes = self.mem_manager.alloc_buffer(len(reqs))
+        # req_ids_gpu = req_ids.cuda()
 
         for i, req in enumerate(reqs):
             input_token_ids = req.get_input_token_ids()
@@ -88,16 +96,22 @@ def insert_for_hybrid_radix_cache(self, reqs):
             # 分配新的 buffer 并复制当前 buffer 的内容
             self.mem_manager.copy_buffer(cur_buffer_idx, new_buffer_indexes[i])
 
-            _, new_shared_kv_node = super().insert(key, value)
+            prefix_len, new_shared_kv_node = super().insert(key, value)
             self.dec_node_ref_counter(req.shared_kv_node)
             self.add_node_ref_counter(new_shared_kv_node)
-            self.set_node_buffer_idx(new_shared_kv_node, new_buffer_indexes[i].item())
+            self.add_buffer_idx_to_node(new_shared_kv_node, new_buffer_indexes[i].item())
             req.shared_kv_node = new_shared_kv_node
+            # 更新 prompt_cache_len，这样 free_a_req_mem 不会释放已属于树的 token
+            # free_a_req_mem 中会释放 [prompt_cache_len:prefix_len]，更新后这个范围为空
+            req.shm_req.prompt_cache_len = req.cur_kv_len
 
     def match_prefix(self, key, update_refs=False):
         assert len(key) != 0
+        self.match_count = (self.match_count + 1) % self.log_interval
+        self.match_len += len(key)
         ans_value_list = []
         tree_node = self._match_prefix_helper(self.root_node, key, ans_value_list, update_refs=update_refs)
+        origin_ans_len = sum(len(v) for v in ans_value_list)
         evict_token_list = []
         while tree_node != self.root_node and tree_node.buffer_idx is None:
             if tree_node.is_leaf():
@@ -126,7 +140,7 @@ def match_prefix(self, key, update_refs=False):
             self.mem_manager.free(evict_token_value)
 
         if tree_node == self.root_node:
-            return None, 0, None
+            return None, origin_ans_len, None
 
         update_node = tree_node
         while update_node != self.root_node:
@@ -137,14 +151,31 @@ def match_prefix(self, key, update_refs=False):
             update_node = update_node.parent
 
         value = torch.concat(ans_value_list)
-        return tree_node, len(value), value
-
-    def set_node_buffer_idx(self, node: TreeNode, buffer_idx: int):
+        # logger.info("HybridRadixCache match_prefix hit tokens: {}".format(len(value)))
+        self.hit_len += len(value)
+        if self.match_count == 0:
+            logger.info(
+                f"HybridRadixCache match_prefix avg hit rate: {self.hit_len / self.match_len:.4f} "
+                f"({self.hit_len}/{self.match_len}) over last {self.log_interval} matches"
+            )
+            self.match_len = 0
+            self.hit_len = 0
+
+        return tree_node, origin_ans_len, value
+
+    def add_buffer_idx_to_node(self, node: TreeNode, buffer_idx: int):
         """Set buffer_idx for a node and add it to evict_buffer_set."""
-        node.buffer_idx = buffer_idx
         self.evict_buffer_set.discard(node)
+        if node.is_leaf():
+            self.evict_tree_set.discard(node)
+        if node.buffer_idx is not None:
+            self.mem_manager.free_buffer([node.buffer_idx])
+        node.buffer_idx = buffer_idx
         node.update_buffer_time()
         self.evict_buffer_set.add(node)
+        if node.is_leaf():
+            self.evict_tree_set.add(node)
+        return
 
     def free_radix_cache_to_get_enough_token(self, need_token_num):
         assert self.mem_manager is not None

lightllm/server/router/model_infer/infer_batch.py

@@ -71,39 +71,31 @@ def get_cpu_kv_cache_stream(self) -> torch.cuda.Stream:
             self.cpu_kv_cache_stream = torch.cuda.Stream()
         return self.cpu_kv_cache_stream
 
-    def _maybe_alloc_and_copy_req_buffers(self, req_objs: List["InferReq"]) -> None:
-        """
-        For hybrid/linear-attention models (e.g. Qwen3-Next) we allocate a fixed-size buffer per request.
-        If radix cache hits and the matched node has a buffer, copy that buffer content to the newly
-        allocated buffer for this request.
-        """
-        if not self.use_buffer_manager or not req_objs:
-            return
-
+    def _alloc_and_copy_req_buffers(self, req_objs: List["InferReq"]) -> None:
+        # 为请求分配 buffer， 如果 shared_kv_node 不为 None，则从 radix cache 复制 buffer。
         if self.radix_cache is not None:
-            # Ensure enough buffer capacity by evicting radix cache buffers if needed.
             self.radix_cache.free_radix_cache_to_get_enough_buffer(len(req_objs))
 
-        req_idxs = np.array([r.req_idx for r in req_objs], dtype=np.int64)
-        request_indices_gpu = torch.from_numpy(req_idxs).to(device="cuda", dtype=torch.int64)
+        req_idxs = []
+        copy_indices = []
+        copy_buffers = []
+
+        for r in req_objs:
+            req_idxs.append(r.req_idx)
+            if r.shared_kv_node is not None:
+                copy_indices.append(r.req_idx)
+                copy_buffers.append(r.shared_kv_node.buffer_idx)
+
+        request_indices_gpu = torch.tensor(req_idxs, device="cuda", dtype=torch.int64)
         self.req_manager.alloc_buffer_for_req(request_indices_gpu)
 
         if self.radix_cache is None:
             return
 
-        # `shared_kv_node` may be None on cache miss; treat it as "no buffer to copy".
-        buffer_idxs = np.array(
-            [None if r.shared_kv_node is None else r.shared_kv_node.buffer_idx for r in req_objs], dtype=object
-        )
-        mask = buffer_idxs == None  # noqa: E711 (intentional elementwise comparison against None)
-        copy_indices = req_idxs[~mask].tolist()
-        if not copy_indices:
-            return
-
-        copy_buffers = buffer_idxs[~mask].tolist()
-        copy_indices_tensor = torch.tensor(copy_indices, device="cuda", dtype=torch.int64)
-        copy_buffers_tensor = torch.tensor(copy_buffers, device="cuda", dtype=torch.int64)
-        self.req_manager.copy_buffer_from_another_buffer(copy_buffers_tensor, copy_indices_tensor)
+        if copy_indices:
+            copy_indices_tensor = torch.tensor(copy_indices, device="cuda", dtype=torch.int64)
+            copy_buffers_tensor = torch.tensor(copy_buffers, device="cuda", dtype=torch.int64)
+            self.req_manager.copy_buffer_from_another_buffer(copy_buffers_tensor, copy_indices_tensor)
 
     def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache: bool = True) -> List["InferReq"]:
         req_objs = []
@@ -143,8 +135,8 @@ def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache:
                     slave_req: InferReq = slave_req
                     slave_req.related_master_req = master_req
 
-        # Hybrid/linear-attention models
-        self._maybe_alloc_and_copy_req_buffers(req_objs)
+        if self.use_buffer_manager and len(req_objs) > 0:
+            self._alloc_and_copy_req_buffers(req_objs)
 
         return req_objs
 
@@ -169,11 +161,11 @@ def free_a_req_mem(self, free_token_index: List, req: "InferReq", free_buffer_in
             if self.use_buffer_manager:
                 buffer_idx = self.req_manager.req_to_buffer_index[req.req_idx].item()
                 if node.buffer_idx is None:
-                    self.radix_cache.set_node_buffer_idx(node, buffer_idx)
+                    self.radix_cache.add_buffer_idx_to_node(node, buffer_idx)
                 else:
                     free_buffer_index.append(buffer_idx)
 
-            old_prefix_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len
+            old_prefix_len = req.shm_req.prompt_cache_len
             free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][old_prefix_len:prefix_len])
             if req.shared_kv_node is not None:
                 assert req.shared_kv_node.node_prefix_total_len <= prefix_len
@@ -218,7 +210,6 @@ def _filter(self, finished_request_ids: List[int]):
             self.req_manager.free(free_req_index, free_token_index)
 
         if self.use_buffer_manager and len(free_buffer_index) != 0:
-            free_buffer_index = torch.tensor(free_buffer_index, dtype=torch.int64, device="cpu")
             self.req_manager.free_buffer(free_buffer_index)
 
         finished_req_ids_set = set(finished_request_ids)
@@ -278,6 +269,7 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool):
     def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bool, can_alloc_token_num: int):
         if paused_reqs:
             g_infer_state_lock.acquire()
+            revovered_reqs = []
             for req in paused_reqs:
                 prefill_need_token_num = req.get_cur_total_len()
                 if prefill_need_token_num > can_alloc_token_num:
@@ -288,8 +280,10 @@ def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bo
                 if is_master_in_dp:
                     req.shm_req.is_paused = False
                 can_alloc_token_num -= prefill_need_token_num
+                revovered_reqs.append(req)
 
-            self._maybe_alloc_and_copy_req_buffers(paused_reqs)
+            self._alloc_and_copy_req_buffers(revovered_reqs)
+            g_infer_state_lock.release()
         return
 
     def get_can_alloc_token_num(self):
@@ -413,14 +407,13 @@ def __init__(
         self.nixl_pd_task_failed_num: int = 0
         self.nixl_trans_device_id: int = -1
 
+        # 在开启radix cache的情况下，用于标记命中情况，用于插入算法
+        self.mamba_model_match_len = 0
+
         # 在开启 enable_cpu_cache 的情况下，当请求结束后，会将请求的 kv cache
         # 卸载到 cpu cache 中，该标志变量用于标记请求的卸载任务的状态
         self.cpu_cache_task_status: "InferReq._CpuCacheTaskStatus" = InferReq._CpuCacheTaskStatus.NOT_STARTED
 
-        # 用于管理该请求整个生命周期固定大小的 buffer 索引，None 表示未分配
-        # 用于线性注意力模型，比如 Qwen3-Next
-        self.buffer_idx: int = None
-
         # mtp_step 用来记录一个请求 draft模型每步需要生成的token数量
         # 正常模式下，这个值为0，在 mtp 模式下，这个值为 draft 模型每步需要生成的token数量
         self.mtp_step: int = get_env_start_args().mtp_step
@@ -469,6 +462,7 @@ def _match_radix_cache(self):
             key = torch.tensor(input_token_ids, dtype=torch.int64, device="cpu")
             key = key[0 : len(key) - 1]  # 最后一个不需要，因为需要一个额外的token，让其在prefill的时候输出下一个token的值
             share_node, kv_len, value_tensor = g_infer_context.radix_cache.match_prefix(key, update_refs=True)
+            self.mamba_model_match_len = kv_len
             if share_node is not None:
                 self.shared_kv_node = share_node
                 ready_cache_len = share_node.node_prefix_total_len