fix

Author: hiworldwzj

lightllm/server/api_cli.py

@@ -189,7 +189,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--router_max_wait_tokens",
         type=int,
-        default=6,
+        default=1,
         help="schedule new requests after every router_max_wait_tokens decode steps.",
     )
     parser.add_argument(

lightllm/server/core/objs/start_args_type.py

@@ -42,7 +42,7 @@ class StartArgs:
     log_stats_interval: int = field(default=10)
     router_token_ratio: float = field(default=0.0)
     router_max_new_token_len: int = field(default=1024)
-    router_max_wait_tokens: int = field(default=6)
+    router_max_wait_tokens: int = field(default=1)
     dp_prefill_wait_step: int = field(default=0)
     disable_aggressive_schedule: bool = field(default=False)
     disable_dynamic_prompt_cache: bool = field(default=False)

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

@@ -0,0 +1,73 @@
+from enum import Enum
+from typing import List
+from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.server.router.model_infer.infer_batch import InferReq
+
+class ControlState:
+
+    def __init__(self):
+        self.is_aggressive_schedule = not get_env_start_args().disable_aggressive_schedule
+
+        # 非激进调度参数
+        self.decode_max_step = max(1, get_env_start_args().router_max_wait_tokens)
+        self.left_decode_num = self.decode_max_step
+
+        self.step_count = 0
+
+
+    def select_run_way(self, prefill_reqs: List[InferReq], decode_reqs: List[InferReq]) -> "RunWay":
+        """
+        判断决策运行方式：
+        返回值: RunWay
+        """
+        self.step_count += 1
+        if self.is_aggressive_schedule:
+            return self._agressive_way(prefill_reqs=prefill_reqs,
+                                       decode_reqs=decode_reqs)
+        else:
+            return self._normal_way(prefill_reqs=prefill_reqs,
+                                    decode_reqs=decode_reqs)
+
+    def _agressive_way(self, prefill_reqs: List[InferReq], decode_reqs: List[InferReq]):
+        if prefill_reqs:
+            return RunWay.PREFILL
+        if decode_reqs:
+            return RunWay.DECODE
+        return RunWay.PASS
+    
+    def _normal_way(self, prefill_reqs: List[InferReq], decode_reqs: List[InferReq]):
+        if decode_reqs:
+            if self.left_decode_num > 0:
+                self.left_decode_num -= 1
+                return RunWay.DECODE
+            else:
+                if prefill_reqs:
+                    self.left_decode_num = self.decode_max_step
+                    return RunWay.PREFILL
+                else:
+                    return RunWay.DECODE
+        else:
+            if prefill_reqs:
+                self.left_decode_num = self.decode_max_step
+                return RunWay.PREFILL
+            else:
+                return RunWay.PASS
+            
+    def try_recover_paused_reqs(self) -> bool:
+        return self.step_count % 100 == 0
+
+        
+
+class RunWay(Enum):
+    PREFILL = 1
+    DECODE = 2
+    PASS = 3
+
+    def is_prefill(self):
+        return self == RunWay.PREFILL
+
+    def is_decode(self):
+        return self == RunWay.DECODE
+
+    def is_pass(self):
+        return self == RunWay.PASS

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

@@ -17,13 +17,17 @@
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.dist_utils import get_current_device_id
 from lightllm.utils.envs_utils import get_env_start_args
+from .control_state import ControlState
 
 logger = init_logger(__name__)
 
 
 class ChunkedPrefillBackend(ModeBackend):
     def __init__(self) -> None:
         super().__init__()
+        
+        # 用于控制每一步是执行prefill 和 decode 还是跳过
+        self.control_state_machine = ControlState()
 
         # 在 mtp 模式下切换绑定的prefill 和 decode 函数
         if get_env_start_args().mtp_mode:
@@ -43,25 +47,29 @@ def infer_loop(self):
 
                 self._try_read_new_reqs()
 
-                prefill_reqs, decode_reqs = self._get_classed_reqs()
-                if prefill_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)
+                
+                if run_way.is_prefill():
                     self.prefill(
                         event_pack=event_pack,
                         prefill_reqs=prefill_reqs,
                     )
                     continue
-
-                if decode_reqs:
+                elif run_way.is_decode():
                     self.decode(
                         event_pack=event_pack,
                         decode_reqs=decode_reqs,
                     )
                     continue
-
-                event_pack.notify_post_handle_and_wait_pre_post_handle()
-                event_pack.notify_forward_and_wait_post_handle()
-                event_pack.notify_pre_post_handle()
-                continue
+                elif run_way.is_pass():
+                    event_pack.notify_post_handle_and_wait_pre_post_handle()
+                    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