[Bug Fix] Remove follow_bootstrap_room fast path in PD disaggregation DP rank resolution (#22901)

python/sglang/srt/disaggregation/common/conn.py

@@ -443,11 +443,28 @@ class CommonKVSender(BaseKVSender):
             return
 
         self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping)
-        if (
+        if self.kv_mgr.server_args.dp_size > 1:
-            self.kv_mgr.server_args.dp_size > 1
+            if self.kv_mgr.server_args.load_balance_method != "follow_bootstrap_room":
-            and self.kv_mgr.server_args.load_balance_method != "follow_bootstrap_room"
+                self._register_prefill_dp_rank()
-        ):
+            elif (
-            self._register_prefill_dp_rank()
+                self.kv_mgr.attn_dp_rank
+                != self.bootstrap_room % self.kv_mgr.server_args.dp_size
+            ):
+                # follow_bootstrap_room was overridden by external routed_dp_rank
+                if envs.SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK.get():
+                    self._register_prefill_dp_rank()
+                else:
+                    self.kv_mgr.record_failure(
+                        self.bootstrap_room,
+                        f"follow_bootstrap_room conflict: dispatched to dp_rank "
+                        f"{self.kv_mgr.attn_dp_rank} but bootstrap_room "
+                        f"{self.bootstrap_room} implies dp_rank "
+                        f"{self.bootstrap_room % self.kv_mgr.server_args.dp_size}. "
+                        f"Set SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK=1 "
+                        f"to allow mixed routing.",
+                    )
+                    self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
+                    return
 
     def _register_prefill_dp_rank(self):
         """Register this request's prefill dp_rank to the bootstrap server."""

python/sglang/srt/disaggregation/decode.py

@@ -438,7 +438,10 @@ class DecodePreallocQueue:
         if prefill_info.dp_size == 1:
             return 0
 
-        if prefill_info.follow_bootstrap_room:
+        if (
+            prefill_info.follow_bootstrap_room
+            and not envs.SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK.get()
+        ):
             return req.bootstrap_room % prefill_info.dp_size
 
         return None

python/sglang/srt/environ.py

@@ -243,6 +243,7 @@ class Envs:
     SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300)
     SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX")
     SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False)
+    SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK = EnvBool(False)
     # Extra slots in req_to_token_pool for decode workers (only effective when
     # max_num_reqs > 32). Increases pool capacity so more KV cache transfers
     # can overlap with decode execution without raising max_running_requests.