done

Author: sufubao

lightllm/models/qwen3next/model.py

@@ -90,25 +90,9 @@ def _init_mem_manager(self):
 
     @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.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(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
 

lightllm/models/qwen3next/req_manager.py

@@ -31,3 +31,23 @@ def free_buffer(self, free_req_indexes: List[int]):
             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]):
+        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
+
+        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_state_lock.release()
+
+        cur_buffer_indexes[empty_indexes] = new_buffer_indexes
+        self.req_to_buffer_indexes[req_indexes] = cur_buffer_indexes
+        return

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

@@ -110,6 +110,10 @@ def prefill_normal(
         model_input, run_reqs = prepare_prefill_inputs(
             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(