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

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`<br>`--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) |
 

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,12 +508,13 @@ 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.round_robin_counter = (self.round_robin_counter + 1) % len(self.workers)
-                self.workers
+
-            )
+    def follow_bootstrap_room_scheduler(self, req: Req):
-        else:
+        if self.maybe_external_dp_rank_routing(req):
+            return
+
         # Set default bootstrap_room if in FAKE auto mode and room is None
         if (
             req.bootstrap_room is None
@@ -524,30 +525,22 @@ class DataParallelController:
                 self.workers
             )
 
-            assert (
+        assert req.bootstrap_room is not None, (
-                req.bootstrap_room is not None
+            "req.bootstrap_room should not be None. Do not send requests directly to "
-            ), "req.bootstrap_room should not be None. Do not send requests directly to prefill or decode instances, but send to the router instead."
+            "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 decode_round_robin_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)
-            self.round_robin_counter = (self.round_robin_counter + 1) % len(
-                self.workers
-            )
-            return
-        self.round_robin_scheduler(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:
+            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'"
         )
+        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:

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,7 +2220,8 @@ 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`"
+                    "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":
@@ -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",
             ],

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")