fix

Author: hiworldwzj

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

@@ -46,6 +46,10 @@ def __init__(self) -> None:
 
         self.enable_decode_microbatch_overlap = get_env_start_args().enable_decode_microbatch_overlap
         self.enable_prefill_microbatch_overlap = get_env_start_args().enable_prefill_microbatch_overlap
+
+        # 控制分类的参数变量
+        self.classed_req_no_decode = False
+        self.classed_req_strict_prefill = False
         pass
 
     def init_model(self, kvargs):

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

@@ -25,7 +25,7 @@
 class ChunkedPrefillBackend(ModeBackend):
     def __init__(self) -> None:
         super().__init__()
-        
+
         # 用于控制每一步是执行prefill 和 decode 还是跳过
         self.control_state_machine = ControlState()
 
@@ -47,11 +47,14 @@ def infer_loop(self):
 
                 self._try_read_new_reqs()
 
-                prefill_reqs, decode_reqs = self._get_classed_reqs(recover_paused=self.control_state_machine.try_recover_paused_reqs())
-                
-                run_way = self.control_state_machine.select_run_way(prefill_reqs=prefill_reqs,
-                                                             decode_reqs=decode_reqs)
-                
+                prefill_reqs, decode_reqs = self._get_classed_reqs(
+                    no_decode=self.classed_req_no_decode,
+                    strict_prefill=self.classed_req_strict_prefill,
+                    recover_paused=self.control_state_machine.try_recover_paused_reqs(),
+                )
+
+                run_way = self.control_state_machine.select_run_way(prefill_reqs=prefill_reqs, decode_reqs=decode_reqs)
+
                 if run_way.is_prefill():
                     self.prefill(
                         event_pack=event_pack,
@@ -69,7 +72,7 @@ def infer_loop(self):
                     event_pack.notify_forward_and_wait_post_handle()
                     event_pack.notify_pre_post_handle()
                     continue
-                
+
         except BaseException as e:
             self.logger.exception(str(e))
             raise e

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_impl.py

@@ -24,6 +24,7 @@ def __init__(self, info_queue: mp.Queue, mem_queue: mp.Queue) -> None:
         super().__init__()
         self.info_queue: mp.Queue = info_queue
         self.mem_queue: mp.Queue = mem_queue
+        self.classed_req_no_decode = True
 
     def init_custom(self):
 

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_impl_for_dp.py

@@ -13,6 +13,7 @@ def __init__(self, info_queue: mp.Queue, mem_queue: mp.Queue) -> None:
         super().__init__()
         self.info_queue: mp.Queue = info_queue
         self.mem_queue: mp.Queue = mem_queue
+        self.classed_req_no_decode = True
 
     def init_custom(self):
         ChunckedPrefillForPrefillNode.init_custom(self)

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

@@ -59,7 +59,9 @@ def infer_loop(self):
                 self._try_read_new_reqs()
 
                 prefill_reqs, decode_reqs = self._get_classed_reqs(
-                    recover_paused=self.control_state_machine.try_recover_paused_reqs()
+                    no_decode=self.classed_req_no_decode,
+                    strict_prefill=self.classed_req_strict_prefill,
+                    recover_paused=self.control_state_machine.try_recover_paused_reqs(),
                 )
 
                 dp_prefill_req_nums, dp_decode_req_nums = self._dp_all_gather_prefill_and_decode_req_num(