From 366574b2b8fb4183b9ae29100c3f035ed04327e2 Mon Sep 17 00:00:00 2001
From: Shangming Cai <csmthu@gmail.com>
Date: Sat, 28 Feb 2026 16:41:42 +0800
Subject: [PATCH] [PD] Cleanup BootstrapServer init and ready check (#19551)

Signed-off-by: Shangming Cai <csmthu@gmail.com>
---
 python/sglang/srt/disaggregation/base/conn.py |  2 +-
 .../sglang/srt/disaggregation/common/conn.py  | 21 ++++++++++++-------
 python/sglang/srt/managers/disagg_service.py  |  1 -
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py
index 0abc87053..761e72d35 100644
--- a/python/sglang/srt/disaggregation/base/conn.py
+++ b/python/sglang/srt/disaggregation/base/conn.py
@@ -164,4 +164,4 @@ class BaseKVReceiver(ABC):
 
 class BaseKVBootstrapServer(ABC):
     @abstractmethod
-    def __init__(self, host: str, port: int, dp_size: int = 1): ...
+    def __init__(self, host: str, port: int): ...
diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py
index 3569e6805..16cf8b601 100644
--- a/python/sglang/srt/disaggregation/common/conn.py
+++ b/python/sglang/srt/disaggregation/common/conn.py
@@ -128,8 +128,7 @@ class CommonKVManager(BaseKVManager):
             self.is_dummy_cp_rank = (
                 is_mla_backend and self.attn_cp_size > 1 and self.attn_cp_rank != 0
             )
-            if not self.is_dummy_cp_rank:
-                self.register_to_bootstrap()
+            self.register_to_bootstrap()
             self.transfer_infos = {}
             self.decode_kv_args_table = {}
             self.pp_group = get_pp_group()
@@ -663,7 +662,7 @@ class CommonKVReceiver(BaseKVReceiver):
 
 
 class CommonKVBootstrapServer(BaseKVBootstrapServer):
-    def __init__(self, host: str, port: int, dp_size: int = 1):
+    def __init__(self, host: str, port: int):
         self.host = host
         self.port = port
         self.app = web.Application()
@@ -673,7 +672,7 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
         self.pp_size = None
         self.attn_tp_size = None
         self.attn_cp_size = None
-        self.dp_size = dp_size
+        self.dp_size = None
         self.page_size = None
         self.kv_cache_dtype: Optional[str] = None
         self.follow_bootstrap_room: Optional[bool] = None
@@ -694,11 +693,17 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer):
         self.thread.start()
 
     def _is_ready(self) -> bool:
-        if self.attn_tp_size is None or self.pp_size is None:
+        if (
+            self.attn_tp_size is None
+            or self.attn_cp_size is None
+            or self.pp_size is None
+            or self.dp_size is None
+        ):
             return False
-        # TODO: verify this expected count is correct for all parallelism
-        # combinations (CP / DP attention / system DP / TP / PP).
-        expected = self.dp_size * self.attn_tp_size * self.pp_size
+        expected = self.dp_size * self.attn_cp_size * self.attn_tp_size * self.pp_size
+        logger.debug(
+            f"Expected {expected} prefill servers to be registered, {self._registered_count} registered so far"
+        )
         return self._registered_count >= expected
 
     def _setup_routes(self):
diff --git a/python/sglang/srt/managers/disagg_service.py b/python/sglang/srt/managers/disagg_service.py
index e0826d6b8..8b02add20 100644
--- a/python/sglang/srt/managers/disagg_service.py
+++ b/python/sglang/srt/managers/disagg_service.py
@@ -26,7 +26,6 @@ def start_disagg_service(
         bootstrap_server = 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