chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

dp-attention: add follow_bootstrap_room + auto load-balance; drop decode_round_robin (#16110)

Browse Source
This commit is contained in:
Mufeez Amjad
2025-12-30 09:33:06 -05:00
committed by GitHub
parent 4de59d83a1
commit cbff7ad985
4 changed files with 80 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -70,7 +70,7 @@ class LoadBalanceMethod(Enum):
"""Load balance method."""
ROUND_ROBIN = auto()
DECODE_ROUND_ROBIN = auto()
FOLLOW_BOOTSTRAP_ROOM = auto()
SHORTEST_QUEUE = auto()
MINIMUM_TOKENS = auto()
@@ -175,7 +175,7 @@ class DataParallelController:
self.round_robin_counter = 0
dispatch_lookup = {
LoadBalanceMethod.ROUND_ROBIN: self.round_robin_scheduler,
LoadBalanceMethod.DECODE_ROUND_ROBIN: self.decode_round_robin_scheduler,
LoadBalanceMethod.FOLLOW_BOOTSTRAP_ROOM: self.follow_bootstrap_room_scheduler,
LoadBalanceMethod.SHORTEST_QUEUE: self.shortest_queue_scheduler,
LoadBalanceMethod.MINIMUM_TOKENS: self.minimum_tokens_scheduler,
}
@@ -508,46 +508,39 @@ class DataParallelController:
if self.maybe_external_dp_rank_routing(req):
return
if self.server_args.disaggregation_mode == "null":
self.workers[self.round_robin_counter].send_pyobj(req)
self.round_robin_counter = (self.round_robin_counter + 1) % len(
self.workers
)
else:
# Set default bootstrap_room if in FAKE auto mode and room is None
if (
req.bootstrap_room is None
and self.server_args.disaggregation_decode_enable_fake_auto
):
req.bootstrap_room = self.round_robin_counter
self.round_robin_counter = (self.round_robin_counter + 1) % len(
self.workers
)
self.workers[self.round_robin_counter].send_pyobj(req)
self.round_robin_counter = (self.round_robin_counter + 1) % len(self.workers)
assert (
req.bootstrap_room is not None
), "req.bootstrap_room should not be None. Do not send requests directly to prefill or decode instances, but send to the router instead."
target_rank = req.bootstrap_room % len(self.workers)
self.workers[target_rank].send_pyobj(req)
def decode_round_robin_scheduler(self, req: Req):
def follow_bootstrap_room_scheduler(self, req: Req):
if self.maybe_external_dp_rank_routing(req):
return
if self.server_args.disaggregation_mode == "decode":
self.workers[self.round_robin_counter].send_pyobj(req)
# Set default bootstrap_room if in FAKE auto mode and room is None
if (
req.bootstrap_room is None
and self.server_args.disaggregation_decode_enable_fake_auto
):
req.bootstrap_room = self.round_robin_counter
self.round_robin_counter = (self.round_robin_counter + 1) % len(
self.workers
)
return
self.round_robin_scheduler(req)
assert req.bootstrap_room is not None, (
"req.bootstrap_room should not be None. Do not send requests directly to "
"prefill or decode instances; send to the router instead."
)
target_rank = req.bootstrap_room % len(self.workers)
self.workers[target_rank].send_pyobj(req)
def shortest_queue_scheduler(self, req):
if self.maybe_external_dp_rank_routing(req):
return
target_worker = self.dp_budget.dispatch()
if target_worker is None:
self.round_robin_scheduler(req)
if self.server_args.disaggregation_mode == "null":
self.round_robin_scheduler(req)
else:
self.follow_bootstrap_room_scheduler(req)
else:
self.workers[target_worker].send_pyobj(req)
@@ -559,7 +552,10 @@ class DataParallelController:
"The 'minimum_tokens' load balancing method is deprecated for now and will introduced later."
"Fall back to 'round_robin_scheduler'"
)
self.round_robin_scheduler(req)
if self.server_args.disaggregation_mode == "null":
self.round_robin_scheduler(req)
else:
self.follow_bootstrap_room_scheduler(req)
def event_loop(self):
while True: