feat(pd): balance heterogeneous prefill load by input tokens - #1425
Open
sufubao wants to merge 1 commit into
Open
feat(pd): balance heterogeneous prefill load by input tokens#1425sufubao wants to merge 1 commit into
sufubao wants to merge 1 commit into
Conversation
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.
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
背景
bs_balancer这个 DP 均衡器用「请求数」度量每个 dp rank 的负载。在 prefill 节点上，如果 prompt 长短差异大（异构），按请求数均衡会出现：某个 dp rank 被分到一个超长 prompt，其它 rank 按人头看是均衡的、实则近乎空闲，prefill 延迟被拉长。改动
给
DpBsBalancer加一个balance_by_input_tokens模式：负载按「剩余 input token 数」度量，即max(1, input_len - max(0, shm_cur_kv_len))（已命中 KV cache 的部分不计入）。get_dp_balancer在run_mode == "prefill"时启用该模式；decode 节点保持原有按请求数的行为，不受影响。balance_by_input_tokens=False，行为完全等价于改动前。测试
新增
unit_tests/server/router/req_queue/test_dp_bs_balancer.py：test_prefill_balancer_uses_remaining_input_tokens：dp0 已有一个 160k（已缓存 20k → 剩余 140k）的大请求，dp1 有两个 1k 请求；新来一个 10k 请求，按 input-token 均衡应分给 dp1（而不是按人头分给 dp0）。test_decode_balancer_keeps_request_count_behavior：默认（decode）模式下仍按请求数均衡，大请求不影响选点。2 passed。