minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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 01:29:33 -07: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
python/sglang/srt/managers/data_parallel_controller.py
+6 -2
View File
@@ -98,7 +98,7 @@ class DPBudget:
) )
self.total_tokens[load.dp_rank] = load.num_total_tokens 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: if method == LoadBalanceMethod.TOTAL_REQUESTS:
target_rank = self.total_requests.index(min(self.total_requests)) target_rank = self.total_requests.index(min(self.total_requests))
elif method == LoadBalanceMethod.TOTAL_TOKENS: elif method == LoadBalanceMethod.TOTAL_TOKENS:
@@ -112,6 +112,7 @@ class DPBudget:
# Increment the load of that worker by one as a heuristic # Increment the load of that worker by one as a heuristic
self.total_requests[target_rank] += 1 self.total_requests[target_rank] += 1
self.total_tokens[target_rank] += estimated_tokens
return target_rank return target_rank
@@ -584,7 +585,10 @@ class DataParallelController:
def total_tokens_scheduler(self, req: Req): def total_tokens_scheduler(self, req: Req):
if self.maybe_external_dp_rank_routing(req): if self.maybe_external_dp_rank_routing(req):
return 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) self.workers[target_worker].send_pyobj(req)
def event_loop(self): def event_loop(self):