update router

Author: shihaobai

lightllm/server/router/manager.py

@@ -197,10 +197,6 @@ async def wait_to_model_ready(self):
         return
 
     def _get_schedule_time_interval(self):
-        if self.running_batch is None:
-            # 没有运行中的 batch 时，每 10ms 触发一次请求调度
-            return 0.01
-
         # dp 模式，为了更好的配平，需要更长的调度间隔，以便于能收到更多的请求
         return self.schedule_time_interval
 
@@ -370,9 +366,7 @@ def _add_req(self, group_req_indexes: GroupReqIndexes):
 
     def _generate_new_batch(self):
         # 调度的时候需要考虑当前运行的batch，和调度了但是暂时还没有推理的部分请求。
-        new_batch = self.req_queue.generate_new_batch(
-            Batch.merge_two_batch(self.running_batch, self.schedule_new_batch)
-        )
+        new_batch = self.req_queue.generate_new_batch(self.schedule_new_batch)
         self.schedule_new_batch = Batch.merge_two_batch(self.schedule_new_batch, new_batch)
         return
 
@@ -469,7 +463,7 @@ async def _recv_new_reqs_and_schedule(self):
         if self.is_multinode_tp:
             self._multinode_tp_generate_new_batch()
         else:
-            if self._get_paused_req_num() == 0:
+            if self._get_paused_req_num() == 0 and self.shm_reqs_io_buffer.is_empty():
                 self._generate_new_batch()
         return
 

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

@@ -75,6 +75,7 @@ def init_model(self, kvargs):
         self.chunked_prefill_size = self.args.chunked_prefill_size
         self.return_all_prompt_logprobs = self.args.return_all_prompt_logprobs
         self.use_dynamic_prompt_cache = not self.args.disable_dynamic_prompt_cache
+        self.batch_max_tokens = self.args.batch_max_tokens
         self.eos_id: List[int] = kvargs.get("eos_id", [2])
         self.disable_cudagraph = self.args.disable_cudagraph
         self.is_multinode_tp = self.args.nnodes > 1 and self.args.dp == 1
@@ -391,6 +392,7 @@ def _get_classed_reqs(
         # 请求，其逻辑是不适合的。
         pause_max_req_num = 2
         wait_pause_count = 0
+        prefill_tokens = 0
 
         # 因为会使用到 radix cache 和 mem_manager 的计数信息
         # 所以需要加锁保护。
@@ -439,6 +441,11 @@ def _get_classed_reqs(
                         wait_pause_count += 1
             else:
                 token_num = req_obj.prefill_need_token_num(is_chuncked_prefill=not self.disable_chunked_prefill)
+                if prefill_tokens + token_num > self.batch_max_tokens:
+                    # 跳过等下次prefill，避免oom
+                    prefill_tokens = 0
+                    break
+                prefill_tokens += token_num
                 if token_num <= can_alloc_token_num:
                     prefill_reqs.append(req_obj)
                     can_alloc_token_num -= token_num

lightllm/server/router/req_queue/chunked_prefill/impl.py

@@ -69,6 +69,7 @@ def generate_new_batch(self, current_batch: Batch):
         new_batch_first_router_need_tokens = (
             0 if current_batch is None else current_batch.get_batch_decode_need_tokens()[self.dp_index]
         )
+        print(f"new_batch_first_router_need_tokens: {new_batch_first_router_need_tokens}")
 
         self._init_cache_list(current_batch, is_busy)
         can_run_list = []