tmp

Author: sufubao

lightllm/common/req_manager.py

@@ -7,6 +7,7 @@
 from lightllm.common.basemodel.triton_kernel.gen_sampling_params import update_req_to_token_id_counter
 from lightllm.utils.envs_utils import enable_env_vars, get_env_start_args
 from lightllm.utils.config_utils import get_vocab_size
+from lightllm.server.router.dynamic_prompt.hybrid_radix_cache import HybridMemManager
 
 logger = init_logger(__name__)
 

lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py

@@ -251,8 +251,8 @@ def _linear_attn(
         assert isinstance(infer_state.mem_manager, Qwen3NextMemoryManager)
 
         input = input.view(-1, infer_cls.embed_dim_)
-        buffer_idx = infer_state.req_manager.req_to_buffer_indexes[infer_state.b_req_idx]
-        conv_states, ssm_states = infer_state.mem_manager.get_state_cache_buffer(self.layer_idx_)
+        buffer_idx = infer_state.buffer_indexes
+        conv_states, ssm_states = infer_state.mem_manager.get_buffer(self.layer_idx_)
 
         mixed_qkvzba = layer_weight.linear_in_proj.mm(input)
         q, k, v, z, b, a = self._fix_query_key_value_ba_ordering(mixed_qkvzba)

lightllm/models/qwen3next/mem_manager.py

@@ -50,6 +50,7 @@ def __init__(
         conv_state_shape: Tuple[int, ...],
         ssm_state_dtype: torch.dtype,
         ssm_state_shape: Tuple[int, ...],
+        max_req_num: int,
         always_copy=False,
         mem_fraction=0.9,
     ):
@@ -80,8 +81,6 @@ def __init__(
             f"Ssm state use : "
             f"{self.ssm_state_mem_manager.get_cell_size() * linear_attn_cache_size / 1024 ** 3} GB Memory.\n"
         )
-        self.EMPTY_BUFFER_INDEX = -1
-        self.HOLD_BUFFER_INDEX = self.conv_state_mem_manager.HOLD_TOKEN_MEMINDEX
         super().__init__(full_attn_cache_size, dtype, num_kv_heads, head_dim, layer_num, always_copy, mem_fraction)
 
     @override

lightllm/models/qwen3next/model.py

@@ -11,6 +11,7 @@
 from lightllm.models.qwen3next.mem_manager import Qwen3NextMemoryManager
 from lightllm.server.core.objs.start_args_type import StartArgs
 from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput
+from lightllm.models.qwen3next.req_manager import Qwen3NextReqManager
 
 logger = init_logger(__name__)
 
@@ -83,18 +84,42 @@ def _init_mem_manager(self):
                 self.head_linear_k_dim,
                 self.head_linear_v_dim,
             ),
+            max_req_num=self.max_req_num,
             mem_fraction=self.mem_fraction,
         )
 
+    @override
     def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0):
         from lightllm.common.basemodel.infer_lock import g_infer_state_lock
-        from lightllm.common.basemodel.infer_context import g_infer_context
+        from lightllm.server.router.model_infer.infer_batch import g_infer_context
 
         infer_state = super()._create_inferstate(model_input, microbatch_index)
+
+        buffer_indexes = self.req_manager.req_to_buffer_indexes[model_input.b_req_idx]
+        empty_indexes = buffer_indexes == self.req_manager.EMPTY_BUFFER_INDEX
+        num_empty = empty_indexes.sum()
+        if num_empty == 0:
+            return infer_state
+
         g_infer_state_lock.acquire()
         if g_infer_context.radix_cache is not None:
-            g_infer_context.radix_cache.free_radix_cache_to_get_enough_buffer(infer_state.batch_size)
-        buffer_indexes = self.mem_manager.alloc_buffer(infer_state.batch_size)
+            g_infer_context.radix_cache.free_radix_cache_to_get_enough_buffer(num_empty)
+        new_buffer_indexes = self.mem_manager.alloc_buffer(num_empty).cuda()
         g_infer_state_lock.release()
+
+        buffer_indexes[empty_indexes] = new_buffer_indexes
+        self.req_manager.req_to_buffer_indexes[model_input.b_req_idx] = buffer_indexes
         infer_state.buffer_indexes = buffer_indexes
         return infer_state
+
+    @override
+    def _init_req_manager(self):
+        create_max_seq_len = 0
+
+        if self.batch_max_tokens is not None:
+            create_max_seq_len = max(create_max_seq_len, self.batch_max_tokens)
+        if self.max_seq_length is not None:
+            create_max_seq_len = max(create_max_seq_len, self.max_seq_length)
+
+        self.req_manager = Qwen3NextReqManager(self.max_req_num, create_max_seq_len, self.mem_manager)
+        return

lightllm/models/qwen3next/req_manager.py

@@ -0,0 +1,33 @@
+from typing import override, List
+
+import torch
+
+from lightllm.common.req_manager import ReqManager
+from lightllm.models.qwen3next.mem_manager import Qwen3NextMemoryManager
+
+
+class Qwen3NextReqManager(ReqManager):
+    def __init__(self, max_request_num, max_sequence_length, mem_manager: Qwen3NextMemoryManager):
+        super().__init__(max_request_num, max_sequence_length, mem_manager)
+        self.EMPTY_BUFFER_INDEX = -1
+        self.req_to_buffer_indexes = torch.zeros((self.max_request_num + 1), dtype=torch.int32, device="cuda")
+        self.req_to_buffer_indexes[:] = self.EMPTY_BUFFER_INDEX
+
+    @override
+    def free(self, free_req_indexes: List[int], free_token_index):
+        self.free_buffer(free_req_indexes)
+        super().free(free_req_indexes, free_token_index)
+
+    @override
+    def free_all(self):
+        self.req_to_buffer_indexes[:] = self.EMPTY_BUFFER_INDEX
+        super().free_all()
+        return
+
+    def free_buffer(self, free_req_indexes: List[int]):
+        from lightllm.server.router.model_infer.infer_batch import g_infer_context
+
+        if g_infer_context.radix_cache is None:
+            self.mem_manager.free_buffer(self.req_to_buffer_indexes[free_req_indexes])
+        self.req_to_buffer_indexes[free_req_indexes] = self.EMPTY_BUFFER_INDEX
+        return

lightllm/server/router/dynamic_prompt/hybrid_radix_cache.py

@@ -5,7 +5,6 @@
 
 from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache, TreeNode
 from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager
-from lightllm.server.router.model_infer.infer_batch import InferReq
 
 
 class HybridMemManager(MemoryManager):
@@ -69,7 +68,7 @@ def evict_buffer(self, need_evict_buffer_num, evict_buffer_callback, evict_token
                     self._remove_leaf_node(node)
         return
 
-    def insert_for_hybrid_radix_cache(self, reqs: List["InferReq"]):
+    def insert_for_hybrid_radix_cache(self, reqs):
         from lightllm.server.router.model_infer.infer_batch import g_infer_context
         from lightllm.common.basemodel.infer_lock import g_infer_state_lock
 

lightllm/server/router/manager.py

@@ -251,22 +251,13 @@ async def loop_for_fwd(
                         estimated_peak_token_count = self.shared_token_load.get_estimated_peak_token_count(d_i)
                         paused_req_num = self._get_paused_req_num_in_dp_index(dp_index=d_i)
 
-                        # Get hit rate from radix cache if available
-                        hit_rate = 0.0
-                        if self.radix_cache_client is not None:
-                            try:
-                                hit_rate = self.radix_cache_client.get_match_prefix_hit_rate(d_i)
-                            except Exception as e:
-                                logger.warning(f"Failed to get hit rate from radix cache: {e}")
-
                         logger.debug(
                             f"dp_i {d_i} current batch size: {len(self.running_batch.reqs)} \n"
                             f"dp_i {d_i} paused req num: {paused_req_num} \n"
                             f"dp_i {d_i} frozen token num: {frozen_token_num} \n"
                             f"dp_i {d_i} estimated_peak_token_count: {estimated_peak_token_count} \n"
                             f"dp_i {d_i} token used ratio: {token_ratio1} not contain prompt cache tree unrefed token\n"
                             f"dp_i {d_i} token used ratio: {token_ratio2} contain prompt cache tree unrefed token\n"
-                            f"dp_i {d_i} match_prefix hit_rate: {hit_rate:.4f}"
                         )
                         self.metric_client.gauge_set("lightllm_batch_pause_size", paused_req_num)
                 # pd decode mode need to update token_load more frequently

lightllm/server/router/model_infer/infer_batch.py

@@ -183,8 +183,10 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool):
         if pause_reqs:
             g_infer_state_lock.acquire()
 
+            pause_req_ids = []
             free_token_index = []
             for req in pause_reqs:
+                pause_req_ids.append(req.req_id)
                 if self.args.diverse_mode:
                     # 发生暂停的时候，需要清除 diverse 模式下的主从关系
                     req.clear_master_slave_state()
@@ -201,6 +203,9 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool):
                 free_token_index = custom_cat(free_token_index)
                 self.req_manager.free_token(free_token_index)
 
+            if hasattr(self.req_manager, "free_buffer"):
+                self.req_manager.free_buffer(pause_req_ids)
+
             g_infer_state_lock.release()
         return self
 

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -192,7 +192,6 @@ def init_model(self, kvargs):
             shm_req_manager=self.shm_req_manager,
             vocab_size=self.model.vocab_size,
         )
-
         # 初始化 dp 模式使用的通信 tensor, 对于非dp模式，不会使用到
         if self.dp_size > 1:
             self.dp_reduce_tensor = torch.tensor([0], dtype=torch.int32, device="cuda", requires_grad=False)