diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 5bfe2fcdd..b5a2cfcca 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -205,7 +205,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--data-parallel-size`
`--dp-size` | The data parallelism size. | `1` | Type: int | -| `--load-balance-method` | The load balancing strategy for data parallelism. The Minimum Token algorithm can only be used when DP attention is applied. This algorithm performs load balancing based on the real-time token load of the DP workers. | `round_robin` | `round_robin`, `shortest_queue`, `minimum_tokens` | +| `--load-balance-method` | The load balancing strategy for data parallelism. The Minimum Token algorithm can only be used when DP attention is applied. This algorithm performs load balancing based on the real-time token load of the DP workers. | `auto` | `auto`, `round_robin`, `follow_bootstrap_room`, `shortest_queue`, `minimum_tokens` | | `--load-watch-interval` | The interval of load watching in seconds. | `0.1` | Type: float | | `--prefill-round-robin-balance` | Prefill is round robin balanced. This is used to promise decode server can get the correct dp rank. | `False` | bool flag (set to enable) | diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index d4967da1a..4f297a32d 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -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: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2927b04d1..9910f5f2e 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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", ], diff --git a/test/srt/test_server_args.py b/test/srt/test_server_args.py index 1489160ad..04c6ce22e 100644 --- a/test/srt/test_server_args.py +++ b/test/srt/test_server_args.py @@ -25,6 +25,20 @@ class TestPrepareServerArgs(CustomTestCase): ) +class TestLoadBalanceMethod(unittest.TestCase): + def test_non_pd_defaults_to_round_robin(self): + server_args = ServerArgs(model_path="dummy", disaggregation_mode="null") + self.assertEqual(server_args.load_balance_method, "round_robin") + + def test_pd_prefill_defaults_to_follow_bootstrap_room(self): + server_args = ServerArgs(model_path="dummy", disaggregation_mode="prefill") + self.assertEqual(server_args.load_balance_method, "follow_bootstrap_room") + + def test_pd_decode_defaults_to_round_robin(self): + server_args = ServerArgs(model_path="dummy", disaggregation_mode="decode") + self.assertEqual(server_args.load_balance_method, "round_robin") + + class TestPortArgs(unittest.TestCase): @patch("sglang.srt.server_args.is_port_available") @patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")