From bbcfcaeefe3672a79578b9d7d1548a76df93a4d0 Mon Sep 17 00:00:00 2001 From: Lawrence Wu Date: Fri, 10 Jul 2026 21:08:02 -0700 Subject: [PATCH] fix(pd): do not abort when req.disagg_prefill_dp_rank is used (#27546) Co-authored-by: Qiaolin Yu --- python/sglang/srt/disaggregation/base/conn.py | 1 + python/sglang/srt/disaggregation/common/conn.py | 3 ++- python/sglang/srt/disaggregation/fake/conn.py | 1 + python/sglang/srt/disaggregation/mooncake/conn.py | 10 +++++++++- python/sglang/srt/disaggregation/mori/conn.py | 10 +++++++++- python/sglang/srt/disaggregation/nixl/conn.py | 10 +++++++++- python/sglang/srt/disaggregation/prefill.py | 1 + 7 files changed, 32 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 7412ec9f0..2b7976506 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -112,6 +112,7 @@ class BaseKVSender(ABC): bootstrap_room: int, dest_tp_ranks: List[int], pp_rank: int, + req_has_disagg_prefill_dp_rank: bool = False, ): ... @abstractmethod diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index b51ee6290..8062123e3 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -957,6 +957,7 @@ class CommonKVSender(BaseKVSender): bootstrap_room: int, dest_tp_ranks: List[int], pp_rank: int, + req_has_disagg_prefill_dp_rank: bool = False, ): self.kv_mgr = mgr self.bootstrap_room = bootstrap_room @@ -975,7 +976,7 @@ class CommonKVSender(BaseKVSender): return self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) - if self.kv_mgr.server_args.dp_size > 1: + if self.kv_mgr.server_args.dp_size > 1 and not req_has_disagg_prefill_dp_rank: if self.kv_mgr.server_args.load_balance_method != "follow_bootstrap_room": self._register_prefill_dp_rank() elif ( diff --git a/python/sglang/srt/disaggregation/fake/conn.py b/python/sglang/srt/disaggregation/fake/conn.py index 09c5f80ce..b74238359 100644 --- a/python/sglang/srt/disaggregation/fake/conn.py +++ b/python/sglang/srt/disaggregation/fake/conn.py @@ -43,6 +43,7 @@ class FakeKVSender(BaseKVSender): bootstrap_room: int, dest_tp_ranks: List[int], pp_rank: int, + req_has_disagg_prefill_dp_rank: bool = False, ): self.kv_mgr = mgr self.has_sent = False diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 970c6b7e9..200646d09 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1751,8 +1751,16 @@ class MooncakeKVSender(CommonKVSender): bootstrap_room: int, dest_tp_ranks: List[int], pp_rank: int, + req_has_disagg_prefill_dp_rank: bool = False, ): - super().__init__(mgr, bootstrap_addr, bootstrap_room, dest_tp_ranks, pp_rank) + super().__init__( + mgr, + bootstrap_addr, + bootstrap_room, + dest_tp_ranks, + pp_rank, + req_has_disagg_prefill_dp_rank, + ) self.conclude_state = None self.init_time = time.time() self._init_trace_ctx() diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 9cf9a9b65..61835a9f2 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -1392,8 +1392,16 @@ class MoriKVSender(CommonKVSender): bootstrap_room: int, dest_tp_ranks: List[int], pp_rank: int, + req_has_disagg_prefill_dp_rank: bool = False, ): - super().__init__(mgr, bootstrap_addr, bootstrap_room, dest_tp_ranks, pp_rank) + super().__init__( + mgr, + bootstrap_addr, + bootstrap_room, + dest_tp_ranks, + pp_rank, + req_has_disagg_prefill_dp_rank, + ) self.transfer_statuses: List[TransferStatus] = [] self.pending_infos: Optional[List[TransferInfo]] = None self.conclude_state: Optional[KVPoll] = None diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 5375b7096..a2a1074c4 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -2391,8 +2391,16 @@ class NixlKVSender(CommonKVSender): bootstrap_room: int, dest_tp_ranks: List[int], pp_rank: int, + req_has_disagg_prefill_dp_rank: bool = False, ): - super().__init__(mgr, bootstrap_addr, bootstrap_room, dest_tp_ranks, pp_rank) + super().__init__( + mgr, + bootstrap_addr, + bootstrap_room, + dest_tp_ranks, + pp_rank, + req_has_disagg_prefill_dp_rank, + ) self.has_sent = False self.chunk_id = 0 self._send_failed = False diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index b7e4247c2..1e0fdef8b 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -269,6 +269,7 @@ class PrefillBootstrapQueue: bootstrap_room=req.bootstrap_room, dest_tp_ranks=dest_tp_ranks, pp_rank=self.pp_rank, + req_has_disagg_prefill_dp_rank=req.disagg_prefill_dp_rank is not None, ) self._process_req(req) req.pending_bootstrap = True