Merge remote-tracking branch 'origin' into autotuner2

Author: sufubao

lightllm/common/req_manager.py

@@ -159,6 +159,7 @@ def init_req_sampling_params(self, req):
                 token_id_counter(
                     prompt_ids=prompt_ids, out_token_id_counter=self.req_to_out_token_id_counter[req.req_idx]
                 )
+                torch.cuda.current_stream().synchronize()
 
         return
 

lightllm/server/api_cli.py

@@ -144,6 +144,13 @@ def make_argument_parser() -> argparse.ArgumentParser:
                         using the deepseekv2 model, set dp to be equal to the tp parameter. In other cases, please
                         do not set it and keep the default value as 1.""",
     )
+    parser.add_argument(
+        "--dp_balancer",
+        type=str,
+        default="bs_balancer",
+        choices=["round_robin", "bs_balancer"],
+        help="the dp balancer type, default is bs_balancer",
+    )
     parser.add_argument(
         "--max_req_total_len", type=int, default=16384, help="the max value for req_input_len + req_output_len"
     )

lightllm/server/router/batch.py

@@ -40,6 +40,15 @@ def get_req_list_for_dp(self, dp_index: int):
                 req_list.append(req)
         return req_list
 
+    def get_all_dp_req_num(self) -> List[int]:
+        if self.dp_size_in_node == 1:
+            return [len(self.reqs)]
+
+        all_dp_req_num = [0 for _ in range(self.dp_size_in_node)]
+        for req in self.reqs:
+            all_dp_req_num[req.sample_params.suggested_dp_index] += 1
+        return all_dp_req_num
+
     def filter_out_finished_req(self, shm_req_manager: ShmReqManager):
         unfinished_req_ids = []
         for req in self.reqs:

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
 

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
@@ -395,6 +396,7 @@ def _get_classed_reqs(
         # 请求，其逻辑是不适合的。
         pause_max_req_num = 2
         wait_pause_count = 0
+        prefill_tokens = 0
 
         # 因为会使用到 radix cache 和 mem_manager 的计数信息
         # 所以需要加锁保护。
@@ -443,7 +445,10 @@ 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:
+                    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/base_queue.py

@@ -26,11 +26,6 @@ def __init__(self, args, router, dp_index, dp_size_in_node) -> None:
         self.router_token_ratio = args.router_token_ratio  # ratio to determine whether the router is busy
         self.router_max_new_token_len = args.router_max_new_token_len
 
-    def append(self, req: Req):
-        req.sample_params.suggested_dp_index = self.dp_index
-        self.waiting_req_list.append(req)
-        return
-
     def extend(self, req_group: List[Req]):
         for req in req_group:
             req.sample_params.suggested_dp_index = self.dp_index

lightllm/server/router/req_queue/dp_balancer/__init__.py

@@ -0,0 +1,13 @@
+from .roundrobin import RoundRobinDpBalancer
+from typing import List
+from lightllm.server.router.req_queue.base_queue import BaseQueue
+from .bs import DpBsBalancer
+
+
+def get_dp_balancer(args, dp_size_in_node: int, inner_queues: List[BaseQueue]):
+    if args.dp_balancer == "round_robin":
+        return RoundRobinDpBalancer(dp_size_in_node, inner_queues)
+    elif args.dp_balancer == "bs_balancer":
+        return DpBsBalancer(dp_size_in_node, inner_queues)
+    else:
+        raise ValueError(f"Invalid dp balancer: {args.dp_balancer}")

lightllm/server/router/req_queue/dp_balancer/base.py

@@ -0,0 +1,23 @@
+import random
+from abc import ABC, abstractmethod
+from typing import List, Union
+from lightllm.server.router.req_queue.base_queue import BaseQueue
+from lightllm.server.router.batch import Batch, Req
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class DpBalancer(ABC):
+    """
+    DP负载均衡器基类
+    定义了负载均衡策略的接口，子类可以实现不同的负载均衡算法
+    """
+
+    def __init__(self, dp_size_in_node: int, inner_queues: List[BaseQueue]):
+        self.dp_size_in_node = dp_size_in_node
+        self.inner_queues = inner_queues
+
+    @abstractmethod
+    def assign_reqs_to_dp(self, current_batch: Batch, reqs_waiting_for_dp_index: List[List[Req]]) -> None:
+        pass

lightllm/server/router/req_queue/dp_balancer/bs.py

@@ -0,0 +1,45 @@
+import random
+from typing import List, Union
+from lightllm.server.router.req_queue.base_queue import BaseQueue
+from lightllm.server.router.batch import Batch, Req
+from lightllm.utils.log_utils import init_logger
+from .base import DpBalancer
+
+logger = init_logger(__name__)
+
+
+class DpBsBalancer(DpBalancer):
+    """
+    This balancer is main to balance the batch size of each dp rank.
+    Because, for dp mode, if it exists a dp rank without any request, it will
+    padding a request and cause the waste of GPU compute resource.
+    """
+
+    def __init__(self, dp_size_in_node: int, inner_queues: List[BaseQueue]):
+        super().__init__(dp_size_in_node, inner_queues)
+
+    def assign_reqs_to_dp(self, current_batch: Batch, reqs_waiting_for_dp_index: List[List[Req]]) -> None:
+        if len(reqs_waiting_for_dp_index) == 0:
+            return
+        # calculate the total load of each dp rank
+        all_dp_req_num = [0 for _ in range(self.dp_size_in_node)]
+        if current_batch is not None:
+            all_dp_req_num = current_batch.get_all_dp_req_num()
+        total_load_per_dp = [
+            all_dp_req_num[i] + len(self.inner_queues[i].waiting_req_list) for i in range(self.dp_size_in_node)
+        ]
+        for req_group in reqs_waiting_for_dp_index:
+            # find the dp rank with minimum load
+            min_load = min(total_load_per_dp)
+            select_dp_indexes = [i for i in range(self.dp_size_in_node) if total_load_per_dp[i] == min_load]
+            suggested_dp_index = random.choice(select_dp_indexes)
+
+            # assign the request to the dp rank and update the load count
+            for req in req_group:
+                req.sample_params.suggested_dp_index = suggested_dp_index
+            self.inner_queues[suggested_dp_index].extend(req_group)
+            # update the load count for this dp rank
+            total_load_per_dp[suggested_dp_index] += len(req_group)
+
+        reqs_waiting_for_dp_index.clear()
+        return

lightllm/server/router/req_queue/dp_balancer/roundrobin.py

@@ -0,0 +1,51 @@
+import random
+from typing import List, Union
+from lightllm.server.router.req_queue.base_queue import BaseQueue
+from lightllm.server.router.batch import Batch, Req
+from lightllm.utils.log_utils import init_logger
+from .base import DpBalancer
+
+logger = init_logger(__name__)
+
+
+class RoundRobinDpBalancer(DpBalancer):
+    """
+    轮询负载均衡器
+    在队列长度最小的DP中进行轮询选择
+    """
+
+    def __init__(self, dp_size_in_node: int, inner_queues: List[BaseQueue]):
+        super().__init__(dp_size_in_node, inner_queues)
+        self.pre_select_dp_index = self.dp_size_in_node - 1
+
+    def get_suggest_dp_index(
+        self,
+    ) -> int:
+        min_length = min(len(queue.waiting_req_list) for queue in self.inner_queues)
+        select_dp_indexes = [
+            i for i, queue in enumerate(self.inner_queues) if len(queue.waiting_req_list) == min_length
+        ]
+
+        # 如果没有可选择的索引，随机选择一个
+        if not select_dp_indexes:
+            self.pre_select_dp_index = random.randint(0, self.dp_size_in_node - 1)
+            return self.pre_select_dp_index
+
+        # 轮询选择
+        for i in range(self.dp_size_in_node):
+            next_dp_index = (self.pre_select_dp_index + i + 1) % self.dp_size_in_node
+            if next_dp_index in select_dp_indexes:
+                self.pre_select_dp_index = next_dp_index
+                return self.pre_select_dp_index
+
+        self.pre_select_dp_index = random.choice(select_dp_indexes)
+        return self.pre_select_dp_index
+
+    def assign_reqs_to_dp(self, current_batch: Batch, reqs_waiting_for_dp_index: List[List[Req]]) -> None:
+        for req_group in reqs_waiting_for_dp_index:
+            suggested_dp_index = self.get_suggest_dp_index()
+            for req in req_group:
+                req.sample_params.suggested_dp_index = suggested_dp_index
+            self.inner_queues[suggested_dp_index].extend(req_group)
+        reqs_waiting_for_dp_index.clear()
+        return