chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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:

43
python/sglang/srt/server_args.py
View File

@@ -376,7 +376,7 @@ class ServerArgs:
# Data parallelism
dp_size: int = 1
load_balance_method: str = "round_robin"
load_balance_method: str = "auto"
# FIXME: remove this after dp rank scheduling is fully supported with PD-Disaggregation
prefill_round_robin_balance: bool = False
@@ -650,6 +650,9 @@ class ServerArgs:
Orchestrates the handling of various server arguments, ensuring proper configuration and validation.
"""
# Normalize load balancing defaults early (before dummy-model short-circuit).
self._handle_load_balance_method()
if self.model_path.lower() in ["none", "dummy"]:
# Skip for dummy models
return
@@ -731,6 +734,36 @@ class ServerArgs:
# Handle any other necessary validations.
self._handle_other_validations()
def _handle_load_balance_method(self):
if self.disaggregation_mode not in ("null", "prefill", "decode"):
raise ValueError(
f"Invalid disaggregation_mode={self.disaggregation_mode!r}"
)
if self.load_balance_method == "auto":
# Default behavior:
# - non-PD: round_robin
# - PD prefill: follow_bootstrap_room
# - PD decode: round_robin
self.load_balance_method = (
"follow_bootstrap_room"
if self.disaggregation_mode == "prefill"
else "round_robin"
)
return
# Backward compat: in PD prefill, legacy "round_robin" means `bootstrap_room` routing.
if (
self.disaggregation_mode == "prefill"
and self.load_balance_method == "round_robin"
):
logger.warning(
"In PD-disaggregation prefill mode, the 'round_robin' load balancing method "
"means `bootstrap_room` routing (use 'follow_bootstrap_room' instead). "
"Falling back to 'follow_bootstrap_room' for backward compatibility."
)
self.load_balance_method = "follow_bootstrap_room"
def _handle_deprecated_args(self):
# Handle deprecated tool call parsers
deprecated_tool_call_parsers = {"qwen25": "qwen", "glm45": "glm"}
@@ -2187,8 +2220,9 @@ class ServerArgs:
if self.dp_size > 1 and not is_in_ci():
assert self.prefill_round_robin_balance, (
"Prefill round robin balance is required when dp size > 1. "
"Please make sure that the prefill instance is launched with `--load-balance-method round_robin`"
" and `--prefill-round-robin-balance` is set for decode server."
"Please make sure that the prefill instance is launched with `--load-balance-method auto` "
"or `--load-balance-method follow_bootstrap_room` "
"and `--prefill-round-robin-balance` is set for decode server."
)
elif self.disaggregation_mode == "prefill":
if self.disaggregation_decode_tp is None:
@@ -3177,8 +3211,9 @@ class ServerArgs:
default=ServerArgs.load_balance_method,
help="The load balancing strategy for data parallelism.",
choices=[
"auto",
"round_robin",
"decode_round_robin",
"follow_bootstrap_room",
"shortest_queue",
"minimum_tokens",
],