kvcache-ai/sglang
1
0
Fork 0
mirror of https://github.com/kvcache-ai/sglang.git synced 2026-07-13 02:46:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix: Add token heuristic increment in total_tokens load balancing (#22614)

Browse Source
This commit is contained in:
ybyang
2026-04-21 16:29:33 +08:00
committed by GitHub
parent 0d69012ef8
commit fa8993111d
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/managers/data_parallel_controller.py
View File

@@ -98,7 +98,7 @@ class DPBudget:
)
self.total_tokens[load.dp_rank] = load.num_total_tokens
def dispatch(self, method: LoadBalanceMethod):
def dispatch(self, method: LoadBalanceMethod, estimated_tokens: int = 0):
if method == LoadBalanceMethod.TOTAL_REQUESTS:
target_rank = self.total_requests.index(min(self.total_requests))
elif method == LoadBalanceMethod.TOTAL_TOKENS:
@@ -112,6 +112,7 @@ class DPBudget:
# Increment the load of that worker by one as a heuristic
self.total_requests[target_rank] += 1
self.total_tokens[target_rank] += estimated_tokens
return target_rank
@@ -584,7 +585,10 @@ class DataParallelController:
def total_tokens_scheduler(self, req: Req):
if self.maybe_external_dp_rank_routing(req):
return
target_worker = self.dp_budget.dispatch(LoadBalanceMethod.TOTAL_TOKENS)
estimated_tokens = len(req.input_ids)
target_worker = self.dp_budget.dispatch(
LoadBalanceMethod.TOTAL_TOKENS, estimated_tokens=estimated_tokens
)
self.workers[target_worker].send_pyobj(req)
def event_loop(self):