feat(pd): balance heterogeneous prefill load by input tokens

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

@@ -8,6 +8,10 @@ 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)
+        return DpBsBalancer(
+            dp_size_in_node,
+            inner_queues,
+            balance_by_input_tokens=args.run_mode == "prefill",
+        )
     else:
         raise ValueError(f"Invalid dp balancer: {args.dp_balancer}")

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

@@ -13,20 +13,42 @@ 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.
+
+    When balance_by_input_tokens is set (e.g. on a prefill node), load is
+    measured in remaining input tokens (input_len minus already-cached kv)
+    instead of request count, so a heterogeneous mix of long and short prompts
+    is balanced by prefill compute rather than headcount.
     """
 
-    def __init__(self, dp_size_in_node: int, inner_queues: List[BaseQueue]):
+    def __init__(
+        self,
+        dp_size_in_node: int,
+        inner_queues: List[BaseQueue],
+        balance_by_input_tokens: bool = False,
+    ):
         super().__init__(dp_size_in_node, inner_queues)
+        self.balance_by_input_tokens = balance_by_input_tokens
+
+    def _req_load(self, req: Req) -> int:
+        if not self.balance_by_input_tokens:
+            return 1
+        return max(1, req.input_len - max(0, req.shm_cur_kv_len))
+
+    def _queue_load(self, reqs: List[Req]) -> int:
+        return sum(self._req_load(req) for req in reqs)
 
     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)]
+        current_load_per_dp = [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()
+            current_load_per_dp = [
+                self._queue_load(current_batch.get_req_list_for_dp(i)) for i in range(self.dp_size_in_node)
+            ]
         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)
+            current_load_per_dp[i] + self._queue_load(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
@@ -39,7 +61,7 @@ def assign_reqs_to_dp(self, current_batch: Batch, reqs_waiting_for_dp_index: Lis
                 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)
+            total_load_per_dp[suggested_dp_index] += self._queue_load(req_group)
 
         reqs_waiting_for_dp_index.clear()
         return

unit_tests/server/router/req_queue/test_dp_bs_balancer.py

@@ -0,0 +1,64 @@
+from types import SimpleNamespace
+
+from lightllm.server.router.req_queue.dp_balancer.bs import DpBsBalancer
+
+
+class FakeReq:
+    def __init__(self, input_len, dp_index, shm_cur_kv_len=0):
+        self.input_len = input_len
+        self.shm_cur_kv_len = shm_cur_kv_len
+        self.sample_params = SimpleNamespace(suggested_dp_index=dp_index)
+
+
+class FakeQueue:
+    def __init__(self):
+        self.waiting_req_list = []
+
+    def extend(self, reqs):
+        self.waiting_req_list.extend(reqs)
+
+
+class FakeBatch:
+    def __init__(self, reqs):
+        self.reqs = reqs
+
+    def get_req_list_for_dp(self, dp_index):
+        return [req for req in self.reqs if req.sample_params.suggested_dp_index == dp_index]
+
+
+def test_prefill_balancer_uses_remaining_input_tokens():
+    queues = [FakeQueue(), FakeQueue()]
+    current_batch = FakeBatch(
+        [
+            FakeReq(160_000, 0, shm_cur_kv_len=20_000),
+            FakeReq(1_000, 1),
+            FakeReq(1_000, 1),
+        ]
+    )
+    new_group = [FakeReq(10_000, -1)]
+
+    balancer = DpBsBalancer(2, queues, balance_by_input_tokens=True)
+    waiting_groups = [new_group]
+    balancer.assign_reqs_to_dp(current_batch, waiting_groups)
+
+    assert new_group[0].sample_params.suggested_dp_index == 1
+    assert queues[1].waiting_req_list == new_group
+    assert waiting_groups == []
+
+
+def test_decode_balancer_keeps_request_count_behavior():
+    queues = [FakeQueue(), FakeQueue()]
+    current_batch = FakeBatch(
+        [
+            FakeReq(160_000, 0),
+            FakeReq(1_000, 1),
+            FakeReq(1_000, 1),
+        ]
+    )
+    new_group = [FakeReq(10_000, -1)]
+
+    balancer = DpBsBalancer(2, queues)
+    balancer.assign_reqs_to_dp(current_batch, [new_group])
+
+    assert new_group[0].sample_params.suggested_dp_index == 0
+    assert queues[0].waiting_req_list == new_group