[PD-Disagg] Support query dp rank from bootstrap server. (#19168)

Signed-off-by: Chang Huaixin (OpenAnolis) <changhuaixin@linux.alibaba.com>
Co-authored-by: Chang Huaixin (OpenAnolis) <changhuaixin@linux.alibaba.com>

docs/advanced_features/pd_disaggregation.md

@@ -142,6 +142,7 @@ The `SGLANG_MOONCAKE_CUSTOM_MEM_POOL` environment variable enables the custom me
 | **`SGLANG_DISAGGREGATION_THREAD_POOL_SIZE`** | Controls the total number of worker threads for KVCache transfer operations per TP rank | A dynamic value calculated by `int(0.75 * os.cpu_count()) // 8)`, which is limited to be larger than 4 and less than 12 to ensure efficiency and prevent thread race conditions |
 | **`SGLANG_DISAGGREGATION_QUEUE_SIZE`** | Sets the number of parallel transfer queues. KVCache transfer requests from multiple decode instances will be sharded into these queues so that they can share the threads and the transfer bandwidth at the same time. If it is set to `1`, then we transfer requests one by one according to fcfs strategy | `4` |
 | **`SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT`** | Timeout (seconds) for receiving destination KV indices during request initialization | `300` |
+| **`SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL`** | Interval (seconds) between cleanups of bootstrap entries | `120` |
 
 If a greater mean TTFT is acceptable, you can `export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600` (10 minutes) to relax the timeout condition.
 Please be aware that this setting will cause prefill instances to take a longer time to clean up the affected memory resources when a running decode node loses connection.

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

@@ -61,6 +61,11 @@ class BaseKVManager(ABC):
         is_mla_backend: Optional[bool] = False,
     ): ...
 
+    @abstractmethod
+    def register_to_bootstrap(self):
+        """Register to the bootstrap server."""
+        ...
+
 
 class BaseKVSender(ABC):
 
@@ -158,4 +163,4 @@ class BaseKVReceiver(ABC):
 
 class BaseKVBootstrapServer(ABC):
     @abstractmethod
-    def __init__(self, host: str, port: int): ...
+    def __init__(self, host: str, port: int, dp_size: int = 1): ...

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

@@ -4,6 +4,7 @@ import asyncio
 import logging
 import socket
 import threading
+import time
 from functools import cache
 from typing import Dict, List, Optional, Tuple, Union
 
@@ -23,6 +24,7 @@ from sglang.srt.disaggregation.base.conn import (
 )
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.distributed import get_pp_group
+from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import (
     get_attention_dp_rank,
     get_attention_dp_size,
@@ -52,6 +54,7 @@ class CommonKVManager(BaseKVManager):
         self.kv_args = args
         self.is_mla_backend = is_mla_backend
         self.disaggregation_mode = disaggregation_mode
+        self.server_args = server_args
         # for p/d multi node infer
         self.bootstrap_host = server_args.host
         self.bootstrap_port = server_args.disaggregation_bootstrap_port
@@ -81,7 +84,7 @@ class CommonKVManager(BaseKVManager):
         self.request_status: Dict[int, KVPoll] = {}
 
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
-            self._register_to_bootstrap()
+            self.register_to_bootstrap()
             self.transfer_infos = {}
             self.decode_kv_args_table = {}
             self.pp_group = get_pp_group()
@@ -93,12 +96,41 @@ class CommonKVManager(BaseKVManager):
             self.prefill_dp_size_table: Dict[str, int] = {}
             self.prefill_pp_size_table: Dict[str, int] = {}
             self.prefill_page_size_table: Dict[str, Optional[int]] = {}
+            self.follow_bootstrap_room_table: Dict[str, bool] = {}
         else:
             raise ValueError(
                 f"Unsupported DisaggregationMode: {self.disaggregation_mode}"
             )
 
-    def _register_to_bootstrap(self):
+    def ensure_parallel_info(self, bootstrap_addr: str) -> bool:
+        """Fetch and cache prefill parallel info if not yet available.
+        Returns True if info is available (cached or freshly fetched).
+        """
+        if bootstrap_addr in self.prefill_dp_size_table:
+            return True
+        info = CommonKVReceiver._fetch_prefill_parallel_info(bootstrap_addr)
+        if info is None:
+            return False
+        tp_size, dp_size, pp_size, page_size, follow_bootstrap_room = info
+
+        if page_size is not None and page_size != self.kv_args.page_size:
+            raise RuntimeError(
+                f"Page size mismatch: prefill server has page_size={page_size}, "
+                f"but decode server has page_size={self.kv_args.page_size}. "
+                f"Both servers must use the same --page-size value."
+            )
+
+        self.prefill_attn_tp_size_table[bootstrap_addr] = tp_size
+        self.prefill_dp_size_table[bootstrap_addr] = dp_size
+        self.prefill_pp_size_table[bootstrap_addr] = pp_size
+        self.prefill_page_size_table[bootstrap_addr] = page_size
+        self.follow_bootstrap_room_table[bootstrap_addr] = follow_bootstrap_room
+        logger.debug(
+            f"Prefill parallel info for [{bootstrap_addr}]: DP={dp_size} TP={tp_size} PP={pp_size} page_size={page_size} follow_bootstrap_room={follow_bootstrap_room}"
+        )
+        return True
+
+    def register_to_bootstrap(self):
         """Register KVSender to bootstrap server via HTTP POST."""
         if self.dist_init_addr:
             # Multi-node case: bootstrap server's host is dist_init_addr
@@ -129,6 +161,7 @@ class CommonKVManager(BaseKVManager):
             "rank_ip": self.local_ip,
             "rank_port": self.rank_port,
             "page_size": self.kv_args.page_size,
+            "load_balance_method": self.server_args.load_balance_method,
         }
 
         try:
@@ -215,6 +248,27 @@ class CommonKVSender(BaseKVSender):
         # inner state
         self.curr_idx = 0
         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()
+
+    def _register_prefill_dp_rank(self):
+        """Register this request's prefill dp_rank to the bootstrap server."""
+        url = f"http://{self.bootstrap_server_url}/register_dp_rank"
+        payload = {
+            "bootstrap_room": self.bootstrap_room,
+            "dp_rank": self.kv_mgr.attn_dp_rank,
+        }
+        try:
+            response = requests.post(url, json=payload, timeout=5)
+            if response.status_code != 200:
+                logger.error(
+                    f"Failed to register prefill dp_rank: {response.status_code}, {response.text}"
+                )
+        except Exception as e:
+            logger.error(f"Failed to register prefill dp_rank: {e}")
 
     def init(self, num_kv_indices: int, aux_index: Optional[int] = None):
         self.num_kv_indices = num_kv_indices
@@ -252,18 +306,7 @@ class CommonKVReceiver(BaseKVReceiver):
         self.kv_mgr = mgr
         self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping)
 
-        if self.bootstrap_addr not in self.kv_mgr.prefill_dp_size_table:
+        if not self.kv_mgr.ensure_parallel_info(self.bootstrap_addr):
-            (
-                self.prefill_attn_tp_size,
-                self.prefill_dp_size,
-                self.prefill_pp_size,
-                self.prefill_page_size,
-            ) = self._get_prefill_parallel_info_from_server()
-            if (
-                self.prefill_attn_tp_size is None
-                or self.prefill_dp_size is None
-                or self.prefill_pp_size is None
-            ):
             self.kv_mgr.record_failure(
                 self.bootstrap_room,
                 f"Could not fetch prefill parallel info from bootstrap_addr: {self.bootstrap_addr}",
@@ -272,42 +315,11 @@ class CommonKVReceiver(BaseKVReceiver):
             self.bootstrap_infos = None
             return
 
-            if self.prefill_page_size is not None:
-                decode_page_size = self.kv_mgr.kv_args.page_size
-                if self.prefill_page_size != decode_page_size:
-                    error_msg = (
-                        f"Page size mismatch: prefill server has page_size={self.prefill_page_size}, "
-                        f"but decode server has page_size={decode_page_size}. "
-                        f"Both servers must use the same --page-size value."
-                    )
-                    logger.error(error_msg)
-                    raise RuntimeError(error_msg)
-
-            logger.debug(
-                f"Fetch prefill parallel info from [{self.bootstrap_addr}]: DP size:{self.prefill_dp_size}, TP size:{self.prefill_attn_tp_size} PP size:{self.prefill_pp_size} Page size:{self.prefill_page_size}"
-            )
-            self.kv_mgr.prefill_attn_tp_size_table[self.bootstrap_addr] = (
-                self.prefill_attn_tp_size
-            )
-            self.kv_mgr.prefill_dp_size_table[self.bootstrap_addr] = (
-                self.prefill_dp_size
-            )
-            self.kv_mgr.prefill_pp_size_table[self.bootstrap_addr] = (
-                self.prefill_pp_size
-            )
-            self.kv_mgr.prefill_page_size_table[self.bootstrap_addr] = (
-                self.prefill_page_size
-            )
-        else:
         self.prefill_attn_tp_size = self.kv_mgr.prefill_attn_tp_size_table[
             self.bootstrap_addr
         ]
-            self.prefill_dp_size = self.kv_mgr.prefill_dp_size_table[
+        self.prefill_dp_size = self.kv_mgr.prefill_dp_size_table[self.bootstrap_addr]
-                self.bootstrap_addr
+        self.prefill_pp_size = self.kv_mgr.prefill_pp_size_table[self.bootstrap_addr]
-            ]
-            self.prefill_pp_size = self.kv_mgr.prefill_pp_size_table[
-                self.bootstrap_addr
-            ]
         self.prefill_page_size = self.kv_mgr.prefill_page_size_table.get(
             self.bootstrap_addr
         )
@@ -367,14 +379,6 @@ class CommonKVReceiver(BaseKVReceiver):
                     self.prefill_attn_tp_size // self.kv_mgr.attn_tp_size
                 ) * (self.prefill_pp_size // self.kv_mgr.pp_size)
 
-        if prefill_dp_rank is not None:
-            logger.debug(f"Targeting DP rank: {prefill_dp_rank}")
-            self.prefill_dp_rank = prefill_dp_rank
-        else:
-            self.prefill_dp_rank = bootstrap_room % self.prefill_dp_size
-
-        self.target_dp_group = self.prefill_dp_rank
-
         # Decode pp size should be equal to prefill pp size or 1
         assert (
             self.kv_mgr.pp_size == self.prefill_pp_size or self.kv_mgr.pp_size == 1
@@ -389,9 +393,17 @@ class CommonKVReceiver(BaseKVReceiver):
         self.kv_mgr.required_prefill_response_num_table[self.bootstrap_room] = (
             self.required_prefill_response_num
         )
-        # NOTE: key distinguished by bootstrap_addr, target_dp_group, and target_tp_rank
+
+        assert (
+            prefill_dp_rank is not None
+        ), "prefill_dp_rank must be resolved before creating receiver"
+        self.prefill_dp_rank = prefill_dp_rank
+        self._setup_bootstrap_infos()
+
+    def _setup_bootstrap_infos(self):
+        # NOTE: key distinguished by bootstrap_addr, prefill_dp_rank, and target_tp_rank
         bootstrap_key = (
-            f"{self.bootstrap_addr}_{self.target_dp_group}_{self.target_tp_rank}"
+            f"{self.bootstrap_addr}_{self.prefill_dp_rank}_{self.target_tp_rank}"
         )
 
         if bootstrap_key not in self.kv_mgr.connection_pool:
@@ -400,7 +412,7 @@ class CommonKVReceiver(BaseKVReceiver):
                 # Enable higher PP ranks to be bootstrapped earlier to make PP PD requests bootstrap more robust
                 for target_pp_rank in reversed(self.target_pp_ranks):
                     bootstrap_info = self._get_bootstrap_info_from_server(
-                        target_tp_rank, self.target_dp_group, target_pp_rank
+                        target_tp_rank, self.prefill_dp_rank, target_pp_rank
                     )
                     if bootstrap_info is not None:
                         if self.kv_mgr.is_mla_backend:
@@ -413,13 +425,13 @@ class CommonKVReceiver(BaseKVReceiver):
                             # For non-MLA: all target_tp_ranks are selected real ranks
                             bootstrap_info["is_dummy"] = False
                         logger.debug(
-                            f"Fetched bootstrap info: {bootstrap_info} for DP {self.target_dp_group} TP {target_tp_rank} PP {target_pp_rank}"
+                            f"Fetched bootstrap info: {bootstrap_info} for DP {self.prefill_dp_rank} TP {target_tp_rank} PP {target_pp_rank}"
                         )
                         bootstrap_infos.append(bootstrap_info)
                     else:
                         self.kv_mgr.record_failure(
                             self.bootstrap_room,
-                            f"Could not fetch bootstrap info for engine rank: {self.kv_mgr.kv_args.engine_rank} and target_dp_group: {self.target_dp_group} and target_pp_rank {target_pp_rank}",
+                            f"Could not fetch bootstrap info for engine rank: {self.kv_mgr.kv_args.engine_rank} and prefill_dp_rank: {self.prefill_dp_rank} and target_pp_rank {target_pp_rank}",
                         )
                         self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
                         return
@@ -435,11 +447,11 @@ class CommonKVReceiver(BaseKVReceiver):
         assert len(self.bootstrap_infos) > 0
 
     def _get_bootstrap_info_from_server(
-        self, engine_rank, target_dp_group, target_pp_rank
+        self, engine_rank, prefill_dp_rank, target_pp_rank
     ):
         """Fetch the bootstrap info from the bootstrap server."""
         try:
-            url = f"http://{self.bootstrap_addr}/route?engine_rank={engine_rank}&target_dp_group={target_dp_group}&target_pp_rank={target_pp_rank}"
+            url = f"http://{self.bootstrap_addr}/route?engine_rank={engine_rank}&prefill_dp_rank={prefill_dp_rank}&target_pp_rank={target_pp_rank}"
             response = requests.get(url, timeout=5)
             if response.status_code == 200:
                 bootstrap_info = response.json()
@@ -453,29 +465,58 @@ class CommonKVReceiver(BaseKVReceiver):
             logger.error(f"Error fetching prefill info from bootstrap: {e}")
             return None
 
-    def _get_prefill_parallel_info_from_server(
+    @staticmethod
-        self,
+    def _fetch_prefill_parallel_info(
-    ) -> Tuple[Optional[int], Optional[int], Optional[int], Optional[int]]:
+        bootstrap_addr: str,
-        """Fetch the prefill parallel info from the bootstrap server."""
+    ) -> Optional[Tuple[int, int, int, int, bool]]:
+        """Fetch the prefill parallel info from the bootstrap server.
+
+        Returns (attn_tp_size, dp_size, pp_size, page_size, follow_bootstrap_room)
+        or None on failure.
+        """
         try:
-            url = f"http://{self.bootstrap_addr}/route?engine_rank={-1}&target_dp_group={-1}&target_pp_rank={-1}"
+            url = f"http://{bootstrap_addr}/route?engine_rank={-1}&prefill_dp_rank={-1}&target_pp_rank={-1}"
-            response = requests.get(url)
+            response = requests.get(url, timeout=5)
             if response.status_code == 200:
-                prefill_parallel_info = response.json()
+                info = response.json()
                 return (
-                    int(prefill_parallel_info["prefill_attn_tp_size"]),
+                    int(info["prefill_attn_tp_size"]),
-                    int(prefill_parallel_info["prefill_dp_size"]),
+                    int(info["prefill_dp_size"]),
-                    int(prefill_parallel_info["prefill_pp_size"]),
+                    int(info["prefill_pp_size"]),
-                    int(prefill_parallel_info["prefill_page_size"]),
+                    int(info["prefill_page_size"]),
+                    bool(info.get("follow_bootstrap_room", True)),
                 )
             else:
                 logger.error(
                     f"Failed to get prefill parallel info: {response.status_code}, {response.text}"
                 )
-                return None, None, None, None
+                return None
         except Exception as e:
             logger.error(f"Error fetching prefill parallel info from bootstrap: {e}")
-            return None, None, None, None
+            return None
+
+    @staticmethod
+    def query_prefill_dp_ranks(
+        bootstrap_addr: str, bootstrap_rooms: List[int]
+    ) -> Dict[str, int]:
+        """Batch query prefill dp_ranks for given bootstrap_rooms."""
+        try:
+            url = f"http://{bootstrap_addr}/query_dp_ranks"
+            response = requests.post(
+                url,
+                json={"bootstrap_rooms": bootstrap_rooms},
+                timeout=5,
+            )
+            if response.status_code == 200:
+                return response.json()
+            else:
+                logger.error(
+                    f"Failed to query dp_ranks: {response.status_code}, {response.text}"
+                )
+                return {}
+        except Exception as e:
+            logger.error(f"Error querying dp_ranks from bootstrap: {e}")
+            return {}
 
     @classmethod
     def _connect(cls, endpoint: str, is_ipv6: bool = False):
@@ -507,7 +548,7 @@ class CommonKVReceiver(BaseKVReceiver):
 
 
 class CommonKVBootstrapServer(BaseKVBootstrapServer):
-    def __init__(self, host: str, port: int):
+    def __init__(self, host: str, port: int, dp_size: int = 1):
         self.host = host
         self.port = port
         self.app = web.Application()
@@ -516,11 +557,16 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
         self._setup_routes()
         self.pp_size = None
         self.attn_tp_size = None
-        self.dp_size = None
+        self.dp_size = dp_size
         self.page_size = None
+        self.follow_bootstrap_room: Optional[bool] = None
         self.prefill_port_table: Dict[
             int, Dict[int, Dict[int, Dict[str, Union[str, int]]]]
         ] = {}
+        self.room_to_dp_rank: Dict[int, Dict[str, Union[int, float]]] = {}
+        self.entry_cleanup_interval = (
+            envs.SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL.get()
+        )
 
         # Start bootstrap server
         self.thread = threading.Thread(target=self._run_server, daemon=True)
@@ -531,6 +577,8 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
 
     def _setup_routes(self):
         self.app.router.add_route("*", "/route", self._handle_route)
+        self.app.router.add_post("/register_dp_rank", self._handle_register_dp_rank)
+        self.app.router.add_post("/query_dp_ranks", self._handle_query_dp_ranks)
         self.app.router.add_get("/health", self._handle_health_check)
 
     async def _handle_health_check(self, request):
@@ -574,6 +622,12 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
         if self.page_size is None and page_size is not None:
             self.page_size = page_size
 
+        if self.follow_bootstrap_room is None:
+            load_balance_method = data.get(
+                "load_balance_method", "follow_bootstrap_room"
+            )
+            self.follow_bootstrap_room = load_balance_method == "follow_bootstrap_room"
+
         if role == "Prefill":
             if system_dp_size == 1:
                 dp_group = attn_dp_rank
@@ -599,15 +653,14 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
 
     async def _handle_route_get(self, request: web.Request):
         engine_rank = request.query.get("engine_rank")
-        target_dp_group = request.query.get("target_dp_group")
+        prefill_dp_rank = request.query.get("prefill_dp_rank")
         target_pp_rank = request.query.get("target_pp_rank")
-        if not engine_rank or not target_dp_group or not target_pp_rank:
+        if not engine_rank or not prefill_dp_rank or not target_pp_rank:
             return web.Response(text="Missing inputs for bootstrap server.", status=400)
 
-        # Currently we use engine_rank == -1 and target_dp_group == -1 to sync dp size
         if (
             int(engine_rank) == -1
-            and int(target_dp_group) == -1
+            and int(prefill_dp_rank) == -1
             and int(target_pp_rank) == -1
         ):
             prefill_parallel_info = {
@@ -615,12 +668,17 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
                 "prefill_dp_size": self.dp_size,
                 "prefill_pp_size": self.pp_size,
                 "prefill_page_size": self.page_size,
+                "follow_bootstrap_room": (
+                    self.follow_bootstrap_room
+                    if self.follow_bootstrap_room is not None
+                    else True
+                ),
             }
             return web.json_response(prefill_parallel_info, status=200)
 
         # Find corresponding prefill info
         async with self.lock:
-            bootstrap_info = self.prefill_port_table[int(target_dp_group)][
+            bootstrap_info = self.prefill_port_table[int(prefill_dp_rank)][
                 int(engine_rank)
             ][int(target_pp_rank)]
 
@@ -629,12 +687,55 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
         else:
             return web.Response(text="Bootstrap info not Found", status=404)
 
+    async def _handle_register_dp_rank(self, request: web.Request):
+        data = await request.json()
+        bootstrap_room = int(data["bootstrap_room"])
+        dp_rank = int(data["dp_rank"])
+        async with self.lock:
+            self.room_to_dp_rank[bootstrap_room] = {
+                "dp_rank": dp_rank,
+                "timestamp": time.time(),
+            }
+        logger.debug(f"Registered dp_rank={dp_rank} for {bootstrap_room=}")
+        return web.Response(text="OK", status=200)
+
+    async def _handle_query_dp_ranks(self, request: web.Request):
+        data = await request.json()
+        bootstrap_rooms = data["bootstrap_rooms"]
+        result = {}
+        async with self.lock:
+            for room in bootstrap_rooms:
+                room_int = int(room)
+                if room_int in self.room_to_dp_rank:
+                    result[str(room_int)] = self.room_to_dp_rank[room_int]["dp_rank"]
+        return web.json_response(result, status=200)
+
+    async def _cleanup_expired_entries(self):
+        """Remove entries older than cleanup interval from room_to_dp_rank."""
+        while True:
+            await asyncio.sleep(self.entry_cleanup_interval)
+            current_time = time.time()
+            async with self.lock:
+                expired_keys = [
+                    key
+                    for key, value in self.room_to_dp_rank.items()
+                    if current_time - value["timestamp"] > self.entry_cleanup_interval
+                ]
+                for key in expired_keys:
+                    del self.room_to_dp_rank[key]
+            if expired_keys:
+                logger.debug(
+                    f"Cleaned up {len(expired_keys)} expired entries from room_to_dp_rank"
+                )
+
     def _run_server(self):
         try:
             # Event Loop
             self._loop = asyncio.new_event_loop()
             asyncio.set_event_loop(self._loop)
 
+            self._loop.create_task(self._cleanup_expired_entries())
+
             access_log = None
             if logging.getLogger(__name__).getEffectiveLevel() <= logging.DEBUG:
                 access_log = self.app.logger

python/sglang/srt/disaggregation/decode.py

@@ -253,6 +253,7 @@ class DecodePreallocQueue:
         # Queue for requests pending pre-allocation
         self.queue: List[DecodeRequest] = []
         self.retracted_queue: List[Req] = []
+        self.pending_reqs: List[Req] = []
         self.prefill_pp_size = prefill_pp_size
         self.kv_manager = self._init_kv_manager()
 
@@ -345,14 +346,41 @@ class DecodePreallocQueue:
             req.retraction_mb_id = None
             self.retracted_queue.append(req)
         else:
-            # Auto enable FAKE mode if configured
+            dp_rank = self._resolve_dp_rank(req)
+            if dp_rank is None:
+                self.pending_reqs.append(req)
+                return
+            self._create_receiver_and_enqueue(req, dp_rank)
+
+    def _resolve_dp_rank(self, req: Req) -> Optional[int]:
+        if req.data_parallel_rank is not None:
+            return req.data_parallel_rank
+
         if req.bootstrap_host == FAKE_BOOTSTRAP_HOST or (
             req.bootstrap_host is None
             and self.scheduler.server_args.disaggregation_transfer_backend == "fake"
         ):
-                kv_receiver_class = get_kv_class(
+            return 0
-                    TransferBackend.FAKE, KVClassType.RECEIVER
+
+        bootstrap_addr = f"{req.bootstrap_host}:{req.bootstrap_port}"
+
+        if bootstrap_addr not in self.kv_manager.prefill_dp_size_table:
+            return None
+
+        if self.kv_manager.follow_bootstrap_room_table[bootstrap_addr]:
+            return (
+                req.bootstrap_room
+                % self.kv_manager.prefill_dp_size_table[bootstrap_addr]
             )
+
+        return None
+
+    def _create_receiver_and_enqueue(self, req: Req, dp_rank: int) -> None:
+        if req.bootstrap_host == FAKE_BOOTSTRAP_HOST or (
+            req.bootstrap_host is None
+            and self.scheduler.server_args.disaggregation_transfer_backend == "fake"
+        ):
+            kv_receiver_class = get_kv_class(TransferBackend.FAKE, KVClassType.RECEIVER)
         else:
             kv_receiver_class = get_kv_class(
                 self.transfer_backend, KVClassType.RECEIVER
@@ -362,7 +390,7 @@ class DecodePreallocQueue:
             mgr=self.kv_manager,
             bootstrap_addr=f"{req.bootstrap_host}:{req.bootstrap_port}",
             bootstrap_room=req.bootstrap_room,
-                prefill_dp_rank=req.data_parallel_rank,
+            prefill_dp_rank=dp_rank,
         )
 
         req.add_latency(RequestStage.DECODE_PREPARE)
@@ -465,10 +493,56 @@ class DecodePreallocQueue:
             else:
                 raise ValueError(f"Unexpected poll case: {poll}")
 
+    def _resolve_pending_reqs(self) -> None:
+        """Batch-resolve dp_ranks for pending requests and create receivers."""
+        if not self.pending_reqs:
+            return
+
+        bootstrap_addr = f"{self.pending_reqs[0].bootstrap_host}:{self.pending_reqs[0].bootstrap_port}"
+
+        # If a request is following the bootstrap room,
+        # we need get the prefill info before resolving the dp_rank,
+        # which is a conflict with the lazy resolve logic in CommonKVReceiver,
+        # so we need to ensure the parallel info before resolving the dp_rank
+        if not self.kv_manager.ensure_parallel_info(bootstrap_addr):
+            return
+
+        resolved = []
+        need_query = []
+        for req in self.pending_reqs:
+            # NOTE: we need resolve it again because we may ensure the parallel info here
+            dp_rank = self._resolve_dp_rank(req)
+            if dp_rank is not None:
+                resolved.append((req, dp_rank))
+            else:
+                need_query.append(req)
+
+        if need_query:
+            from sglang.srt.disaggregation.common.conn import CommonKVReceiver
+
+            rooms = [req.bootstrap_room for req in need_query]
+            room_to_rank = CommonKVReceiver.query_prefill_dp_ranks(
+                bootstrap_addr, rooms
+            )
+            remaining = []
+            for req in need_query:
+                room_key = str(req.bootstrap_room)
+                if room_key in room_to_rank:
+                    resolved.append((req, int(room_to_rank[room_key])))
+                else:
+                    remaining.append(req)
+            self.pending_reqs = remaining
+        else:
+            self.pending_reqs = []
+
+        for req, dp_rank in resolved:
+            self._create_receiver_and_enqueue(req, dp_rank)
+
     def pop_preallocated(
         self, rids_to_check: Optional[List[str]] = None
     ) -> Tuple[List[DecodeRequest], List[DecodeRequest]]:
         """Pop the preallocated requests from the pending queue (FIFO)."""
+        self._resolve_pending_reqs()
         self._update_handshake_waiters(rids_to_check)
 
         failed_reqs = []
@@ -1086,7 +1160,7 @@ class SchedulerDisaggregationDecodeMixin:
         if self.polling_count % self.polling_interval == 0:
             req_conns, _ = self.disagg_decode_prealloc_queue.pop_preallocated()
             self.disagg_decode_transfer_queue.extend(req_conns)
-            alloc_reqs = (
+            transferred_reqs = (
                 self.disagg_decode_transfer_queue.pop_transferred()
             )  # the requests which kv has arrived
-            self.waiting_queue.extend(alloc_reqs)
+            self.waiting_queue.extend(transferred_reqs)

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

@@ -1124,18 +1124,20 @@ class MooncakeKVManager(CommonKVManager):
             ]
             for k in keys_to_remove:
                 del self.connection_pool[k]
-            if failed_bootstrap_addr in self.prefill_attn_tp_size_table:
-                del self.prefill_attn_tp_size_table[failed_bootstrap_addr]
-            if failed_bootstrap_addr in self.prefill_dp_size_table:
-                del self.prefill_dp_size_table[failed_bootstrap_addr]
-            if failed_bootstrap_addr in self.prefill_pp_size_table:
-                del self.prefill_pp_size_table[failed_bootstrap_addr]
 
             possible_affected_rooms = self.addr_to_rooms_tracker.get(
                 failed_bootstrap_addr, []
             )
-            if failed_bootstrap_addr in self.addr_to_rooms_tracker:
+            keys_to_remove = [
-                del self.addr_to_rooms_tracker[failed_bootstrap_addr]
+                self.prefill_attn_tp_size_table,
+                self.prefill_dp_size_table,
+                self.prefill_pp_size_table,
+                self.follow_bootstrap_room_table,
+                self.addr_to_rooms_tracker,
+            ]
+            for k in keys_to_remove:
+                if failed_bootstrap_addr in k:
+                    del k[failed_bootstrap_addr]
 
         # Report the requests associated with the failed bootstrap addr and mark their status as KVPoll.Failed
         affected_rooms = []

python/sglang/srt/environ.py

@@ -259,6 +259,7 @@ class Envs:
     SGLANG_REQ_WAITING_TIMEOUT = EnvFloat(-1)  # in seconds
     SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH = EnvBool(False)
     SGLANG_REQ_RUNNING_TIMEOUT = EnvFloat(-1)  # in seconds
+    SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL = EnvInt(120)
 
     # Test: pd-disaggregation
     SGLANG_TEST_PD_DISAGG_BACKEND = EnvStr("mooncake")

python/sglang/srt/managers/disagg_service.py

@@ -28,6 +28,7 @@ def start_disagg_service(
         bootstrap_server: BaseKVBootstrapServer = kv_bootstrap_server_class(
             host=server_args.host,
             port=server_args.disaggregation_bootstrap_port,
+            dp_size=server_args.dp_size,
         )
         is_create_store = (
             server_args.node_rank == 0 and transfer_backend == TransferBackend.ASCEND

python/sglang/srt/server_args.py

@@ -814,18 +814,6 @@ class ServerArgs:
             )
             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"}

test/registered/distributed/test_disaggregation_dp_attention.py

@@ -20,6 +20,7 @@ register_cuda_ci(est_time=155, suite="stage-c-test-8-gpu-h20")
 class TestDisaggregationDPAttention(PDDisaggregationServerBase):
     PREFILL_DP_SIZE = 4
     DECODE_DP_SIZE = 4
+    LOAD_BALANCE_METHOD = "auto"
 
     @classmethod
     def setUpClass(cls):
@@ -50,6 +51,8 @@ class TestDisaggregationDPAttention(PDDisaggregationServerBase):
             "--dp",
             str(cls.PREFILL_DP_SIZE),
             "--enable-dp-attention",
+            "--load-balance-method",
+            cls.LOAD_BALANCE_METHOD,
         ]
         prefill_args += cls.transfer_backend + cls.rdma_devices
         cls.process_prefill = popen_launch_pd_server(
@@ -72,6 +75,8 @@ class TestDisaggregationDPAttention(PDDisaggregationServerBase):
             "--enable-dp-attention",
             "--base-gpu-id",
             str(cls.PREFILL_DP_SIZE),
+            "--load-balance-method",
+            cls.LOAD_BALANCE_METHOD,
         ]
         decode_args += cls.transfer_backend + cls.rdma_devices
         cls.process_decode = popen_launch_pd_server(
@@ -97,5 +102,9 @@ class TestDisaggregationDPAttention(PDDisaggregationServerBase):
         self.assertGreater(metrics["accuracy"], 0.60)
 
 
+class TestDisaggregationDPAttentionRoundRobin(TestDisaggregationDPAttention):
+    LOAD_BALANCE_METHOD = "round_robin"
+
+
 if __name__ == "__main__":
     unittest.main()