Merge branch 'dp_balancer' of https://github.com/ModelTC/lightllm into dp_balancer

Author: root

lightllm/server/api_cli.py

@@ -119,7 +119,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
         help="tool call parser type",
     )
     parser.add_argument(
-        "--running_max_req_size", type=int, default=2048, help="the max size for forward requests in the same time"
+        "--running_max_req_size", type=int, default=1000, help="the max size for forward requests in the same time"
     )
     parser.add_argument("--nnodes", type=int, default=1, help="the number of nodes")
     parser.add_argument("--node_rank", type=int, default=0, help="the rank of the current node")

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

@@ -446,11 +446,9 @@ def _get_classed_reqs(
             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
+                    continue
                 if token_num <= can_alloc_token_num:
+                    prefill_tokens += token_num
                     prefill_reqs.append(req_obj)
                     can_alloc_token_num -= token_num
                 else:

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

@@ -69,7 +69,6 @@ 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 = []