fix

Author: sufubao

lightllm/common/basemodel/infer_struct.py

@@ -88,9 +88,6 @@ def __init__(self):
         self.dp_output_split_sizes: List[List[int]] = None
         self.dp_input_split_sizes: List[List[int]] = None
 
-        # 专门用于管理混合注意力模型的buffer
-        self.buffer_indexes: torch.Tensor = None
-
     def init_some_extra_state(self, model, input_ids: torch.Tensor):
         if self.is_prefill:
             (

lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py

@@ -251,7 +251,7 @@ def _linear_attn(
         assert isinstance(infer_state.mem_manager, Qwen3NextMemoryManager)
 
         input = input.view(-1, infer_cls.embed_dim_)
-        buffer_idx = infer_state.buffer_indexes
+        buffer_idx = infer_state.req_manager.req_to_buffer_indexes[infer_state.b_req_idx]
         conv_states, ssm_states = infer_state.mem_manager.get_buffer(self.layer_idx_)
 
         mixed_qkvzba = layer_weight.linear_in_proj.mm(input)

lightllm/models/qwen3next/mem_manager.py

@@ -121,8 +121,7 @@ def free_buffer(self, free_buffer_indexes: List[int], reset=True):
     @override
     def alloc_buffer(self, need_size):
         # conv_state 和 ssm_state 共享buffer_idx
-        buffer_indexes = self.conv_state_mem_manager.alloc(need_size)
-        return buffer_indexes
+        return self.conv_state_mem_manager.alloc(need_size)
 
     @override
     def get_buffer_can_use_size(self):

lightllm/models/qwen3next/model.py

@@ -88,14 +88,6 @@ def _init_mem_manager(self):
             mem_fraction=self.mem_fraction,
         )
 
-    @override
-    def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0):
-        infer_state = super()._create_inferstate(model_input, microbatch_index)
-
-        buffer_indexes = self.req_manager.req_to_buffer_indexes[model_input.b_req_idx]
-        infer_state.buffer_indexes = buffer_indexes
-        return infer_state
-
     @override
     def _init_req_manager(self):
         create_max_seq_len = 0

lightllm/models/qwen3next/req_manager.py

@@ -24,30 +24,26 @@ def free_all(self):
         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
-
-    def alloc_buffer(self, req_indexes: List[int]):
+    @override
+    def alloc(self):
         from lightllm.common.basemodel.infer_lock import g_infer_state_lock
         from lightllm.server.router.model_infer.infer_batch import g_infer_context
 
-        cur_buffer_indexes = self.req_to_buffer_indexes[req_indexes]
-        empty_indexes = cur_buffer_indexes == self.EMPTY_BUFFER_INDEX
-        num_empty = empty_indexes.sum()
-        if num_empty == 0:
-            return
+        req_index = super().alloc()
 
         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_token(num_empty)
-        new_buffer_indexes = self.mem_manager.alloc_buffer(num_empty).cuda()
+            g_infer_context.radix_cache.free_radix_cache_to_get_enough_buffer(1)
+        new_buffer_index = self.mem_manager.alloc_buffer(1)
+        self.req_to_buffer_indexes[req_index] = new_buffer_index
         g_infer_state_lock.release()
 
-        cur_buffer_indexes[empty_indexes] = new_buffer_indexes
-        self.req_to_buffer_indexes[req_indexes] = cur_buffer_indexes
+        return req_index
+
+    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

@@ -82,14 +82,12 @@ def insert_for_hybrid_radix_cache(self, reqs):
             input_token_ids = req.get_input_token_ids()
             key = torch.tensor(input_token_ids[0 : req.cur_kv_len], dtype=torch.int64, device="cpu")
             value = g_infer_context.req_manager.req_to_token_indexs[req.req_idx][: req.cur_kv_len].cpu()
-            buffer_idx = req.buffer_idx
 
             # 分配新的 buffer 并复制当前 buffer 的内容
-            self.mem_manager.copy_buffer(buffer_idx, new_buffer_indexes[i])
-            req.buffer_idx = new_buffer_indexes[i]
+            self.mem_manager.copy_buffer(req.buffer_idx, new_buffer_indexes[i])
 
             _, new_shared_kv_node = self.insert(key, value)
-            new_shared_kv_node.buffer_idx = buffer_idx
+            new_shared_kv_node.buffer_idx = new_buffer_indexes[i]
             self.dec_node_ref_counter(req.shared_kv_node)
             self.add_node_ref_counter(new_shared_kv_node)
             req.shared_kv_node = new_shared_kv_node

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

@@ -111,9 +111,6 @@ def prefill_normal(
             prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill, is_multimodal=self.is_multimodal
         )
 
-        if hasattr(g_infer_context.req_manager, "req_to_buffer_indexes"):
-            g_infer_context.req_manager.alloc_buffer(model_input.b_req_idx)
-
         with torch.cuda.stream(g_infer_context.get_overlap_stream()):
             model_output = self.model.forward(model_input)
             _, next_token_ids_cpu, next_token_logprobs_cpu = self._sample_and_scatter_token(
@@ -132,6 +129,9 @@ def prefill_normal(
         event_pack.notify_post_handle_and_wait_pre_post_handle()
         update_packs = self._pre_post_handle(run_reqs, is_chuncked_mode=not self.disable_chunked_prefill)
 
+        if isinstance(g_infer_context.radix_cache, HybridRadixCache):
+            g_infer_context.radix_cache.insert_for_hybrid_radix_cache(run_reqs)
+
         # 第三阶段
         event_pack.notify_forward_and_wait_post_handle()
         sync_event.synchronize()
@@ -143,10 +143,6 @@ def prefill_normal(
             extra_post_req_handle_func=self.extra_post_req_handle_func,
             nixl_prefill_chuncked_handle_func=self.nixl_prefill_chuncked_handle_func,
         )
-
-        if isinstance(g_infer_context.radix_cache, HybridRadixCache):
-            g_infer_context.radix_cache.insert_for_hybrid_radix_cache(run_reqs)
-
         # 第四阶段
         event_pack.notify_pre_post_handle()
         return