[PD] Cleanup BootstrapServer init and ready check (#19551)

Signed-off-by: Shangming Cai <csmthu@gmail.com>

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): ...

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

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