From d9586b0ab5df619257fa492777050e2d12fa1f1f Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 4 Aug 2026 17:29:42 +0800 Subject: [PATCH] feat(pd): balance heterogeneous prefill load by input tokens The bs dp balancer measures load by request count. On a prefill node that is fine when prompts are similar size, but with a heterogeneous mix one dp rank can be assigned a very long prompt while the others sit near-idle headcount-balanced, stretching prefill latency. Add a balance_by_input_tokens mode that measures load as remaining input tokens (input_len minus already-cached kv length). get_dp_balancer enables it for run_mode == "prefill"; decode keeps the request-count behavior unchanged. --- .../router/req_queue/dp_balancer/__init__.py | 6 +- .../server/router/req_queue/dp_balancer/bs.py | 32 ++++++++-- .../router/req_queue/test_dp_bs_balancer.py | 64 +++++++++++++++++++ 3 files changed, 96 insertions(+), 6 deletions(-) create mode 100644 unit_tests/server/router/req_queue/test_dp_bs_balancer.py diff --git a/lightllm/server/router/req_queue/dp_balancer/__init__.py b/lightllm/server/router/req_queue/dp_balancer/__init__.py index 34f994f8a2..a40ca87ff2 100644 --- a/lightllm/server/router/req_queue/dp_balancer/__init__.py +++ b/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}") diff --git a/lightllm/server/router/req_queue/dp_balancer/bs.py b/lightllm/server/router/req_queue/dp_balancer/bs.py index a1fc2195b5..025041a2b4 100644 --- a/lightllm/server/router/req_queue/dp_balancer/bs.py +++ b/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 diff --git a/unit_tests/server/router/req_queue/test_dp_bs_balancer.py b/unit_tests/server/router/req_queue/test_dp_bs_balancer.py new file mode 100644 index 0000000000..3d7791b495 --- /dev/null +++ b/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