From 3600465e81b71523626a582e844626d303013ebe Mon Sep 17 00:00:00 2001
From: Byron Hsu <byronhsu1230@gmail.com>
Date: Wed, 15 Apr 2026 22:53:29 -0700
Subject: [PATCH] [Bug Fix] Remove follow_bootstrap_room fast path in PD
 disaggregation DP rank resolution (#22901)

---
 .../sglang/srt/disaggregation/common/conn.py  | 27 +++++++++++++++----
 python/sglang/srt/disaggregation/decode.py    |  5 +++-
 python/sglang/srt/environ.py                  |  1 +
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py
index 26752d52d..79dfc79c4 100644
--- a/python/sglang/srt/disaggregation/common/conn.py
+++ b/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 (
-            self.kv_mgr.server_args.dp_size > 1
-            and self.kv_mgr.server_args.load_balance_method != "follow_bootstrap_room"
-        ):
-            self._register_prefill_dp_rank()
+        if self.kv_mgr.server_args.dp_size > 1:
+            if self.kv_mgr.server_args.load_balance_method != "follow_bootstrap_room":
+                self._register_prefill_dp_rank()
+            elif (
+                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."""
diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py
index 875900ad8..c55ff312b 100644
--- a/python/sglang/srt/disaggregation/decode.py
+++ b/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
diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index 42b340a3d..88b3a01a9 100644
--- a/python/sglang/srt/environ.py
+++ b/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.