From 5d8dd1f7f99f4c74bc302dd9ef6d32c00c41f7a8 Mon Sep 17 00:00:00 2001 From: ziang663 <1902123738@qq.com> Date: Fri, 24 Jul 2026 03:15:10 +0800 Subject: [PATCH] fix(disagg): honor PP consensus for bootstrap and prealloc --- python/sglang/srt/disaggregation/decode.py | 50 +++++++++++++++---- python/sglang/srt/disaggregation/prefill.py | 46 +++++++++++------ python/sglang/srt/disaggregation/utils.py | 29 ++++++++++- .../sglang/srt/managers/scheduler_pp_mixin.py | 8 +-- .../test_decode_queue_cleanup.py | 2 + ...test_priority_scheduling_disaggregation.py | 1 + .../mem_cache/test_decode_radix_lock_ref.py | 1 + 7 files changed, 106 insertions(+), 31 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 5224920c4481..a664fbe85119 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -55,6 +55,7 @@ is_dsv4_c128_online_enabled, is_mla_backend, poll_and_all_reduce, + poll_and_all_reduce_pp, poll_and_all_reduce_with_staging, prepare_abort, setup_state_kv_args, @@ -325,6 +326,7 @@ def __init__( self.bootstrap_port = bootstrap_port self.max_total_num_tokens = max_total_num_tokens self.pp_rank = pp_rank + self.pp_size = scheduler.ps.pp_size self.num_reserved_decode_tokens = num_reserved_decode_tokens self.transfer_backend = transfer_backend # Queue for requests pending pre-allocation @@ -719,23 +721,40 @@ def resume_retracted_reqs( return resumed_reqs def _update_handshake_waiters( - self, rids_to_check: Optional[List[str]] = None + self, + rids_to_check: Optional[List[str]] = None, + pp_good_rids: Optional[List[str]] = None, + pp_bad_rids: Optional[List[str]] = None, ) -> None: if not self.queue: return # Still poll if any receiver was aborted, otherwise it stays stuck. - if all(decode_req.waiting_for_input for decode_req in self.queue) and not any( - decode_req.kv_receiver.conclude_state == KVPoll.Failed - for decode_req in self.queue + if ( + self.pp_size <= 1 + and all(decode_req.waiting_for_input for decode_req in self.queue) + and not any( + decode_req.kv_receiver.conclude_state == KVPoll.Failed + for decode_req in self.queue + ) ): return - polls = poll_and_all_reduce( - [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group - ) + if self.pp_size > 1: + polls = poll_and_all_reduce_pp( + (decode_req.req.rid for decode_req in self.queue), + KVPoll.WaitingForInput, + pp_good_rids, + pp_bad_rids, + ) + else: + polls = poll_and_all_reduce( + [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group + ) - for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): + for decode_req, poll in zip(self.queue, polls): + if poll is None: + continue if rids_to_check is not None and decode_req.req.rid not in rids_to_check: continue @@ -856,11 +875,22 @@ def _resolve_pending_reqs(self) -> None: decode_req.kv_receiver.init(prefill_dp_rank) def pop_preallocated( - self, rids_to_check: Optional[List[str]] = None + self, + rids_to_check: Optional[List[str]] = None, + pp_good_rids: Optional[List[str]] = None, + pp_bad_rids: Optional[List[str]] = None, ) -> Tuple[List[DecodeRequest], List[DecodeRequest]]: """Pop the preallocated requests from the pending queue (FIFO).""" + is_pp_mode = self.pp_size > 1 + if is_pp_mode and (pp_good_rids is None or pp_bad_rids is None): + raise ValueError("PP consensus is required when pp_size > 1") + if is_pp_mode and rids_to_check is not None: + raise ValueError("rids_to_check cannot be used in PP mode") + self._resolve_pending_reqs() - self._update_handshake_waiters(rids_to_check) + self._update_handshake_waiters(rids_to_check, pp_good_rids, pp_bad_rids) + if is_pp_mode: + rids_to_check = pp_good_rids + pp_bad_rids failed_reqs = [] preallocated_reqs = [] diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 1f4c947fa3fa..51d6e2dba7f2 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -45,6 +45,7 @@ is_dsv4_c128_online_enabled, is_mla_backend, poll_and_all_reduce_attn_cp_tp_group, + poll_and_all_reduce_pp, prepare_abort, setup_state_kv_args, ) @@ -333,13 +334,15 @@ def _process_req(self, req: Req) -> None: def pop_bootstrapped( self, return_failed_reqs: bool = False, - rids_to_check: Optional[List[str]] = None, - ) -> List[Req]: + pp_good_rids: Optional[List[str]] = None, + pp_bad_rids: Optional[List[str]] = None, + ) -> List[Req] | tuple[List[Req], List[Req]]: """ pop the reqs which has finished bootstrapping return_failed_reqs: For PP, on rank 0, also return the failed reqs to notify the next rank - rids_to_check: For PP, on rank > 0, check the rids from the previous rank has consensus with the current rank. + pp_good_rids: RIDs that PP consensus determined as WaitingForInput. + pp_bad_rids: RIDs that PP consensus determined as Failed. """ bootstrapped_reqs = [] @@ -352,21 +355,32 @@ def pop_bootstrapped( else: return [], [] - polls = poll_and_all_reduce_attn_cp_tp_group( - [req.disagg_kv_sender for req in self.queue], - self.scheduler.attn_cp_cpu_group, - self.scheduler.attn_tp_cpu_group, - ) + if self.pp_size > 1: + polls = poll_and_all_reduce_pp( + (req.rid for req in self.queue), + KVPoll.WaitingForInput, + pp_good_rids, + pp_bad_rids, + ) + uncovered = [i for i, poll in enumerate(polls) if poll is None] + if uncovered: + local_polls = poll_and_all_reduce_attn_cp_tp_group( + [self.queue[i].disagg_kv_sender for i in uncovered], + self.scheduler.attn_cp_cpu_group, + self.scheduler.attn_tp_cpu_group, + ) + for i, local_poll in zip(uncovered, local_polls): + if local_poll == KVPoll.Failed: + polls[i] = KVPoll.Failed + else: + polls = poll_and_all_reduce_attn_cp_tp_group( + [req.disagg_kv_sender for req in self.queue], + self.scheduler.attn_cp_cpu_group, + self.scheduler.attn_tp_cpu_group, + ) for i, (req, poll) in enumerate(zip(self.queue, polls)): - if ( - rids_to_check is not None - and req.rid not in rids_to_check - and poll != KVPoll.Failed - ): - # In PP mode, successful bootstrap still requires cross-rank - # consensus. Local failures are terminal and must be drained - # even if an earlier PP rank has already removed the request. + if poll is None: continue if poll == KVPoll.Failed: diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 8d741804d875..7893d588248d 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -5,7 +5,16 @@ from collections import deque from contextlib import nullcontext from enum import Enum -from typing import TYPE_CHECKING, List, Literal, Optional, Tuple, Type, overload +from typing import ( + TYPE_CHECKING, + Iterable, + List, + Literal, + Optional, + Tuple, + Type, + overload, +) import numpy as np import torch @@ -34,6 +43,24 @@ _IS_HIP = is_hip() +def poll_and_all_reduce_pp( + rids: Iterable[str], + ready_poll: int, + pp_good_rids: Optional[List[str]] = None, + pp_bad_rids: Optional[List[str]] = None, +) -> List[Optional[int]]: + """Map authoritative PP consensus to poll states without polling again.""" + if pp_good_rids is None or pp_bad_rids is None: + raise ValueError("PP consensus is required") + + good_rids = set(pp_good_rids) + bad_rids = set(pp_bad_rids) + return [ + KVPoll.Failed if rid in bad_rids else ready_poll if rid in good_rids else None + for rid in rids + ] + + def get_dsa_seed_metadata_dim(hf_config) -> int: """Return the model-defined PD seed width, independent of local spec mode.""" if not getattr(hf_config, "index_share_for_mtp_iteration", False): diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 0a55601d4128..f3e1f76f4ad1 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -811,8 +811,8 @@ def process_bootstrapped_queue( good_reqs, failed_reqs = ( self.disagg_prefill_bootstrap_queue.pop_bootstrapped( return_failed_reqs=True, - rids_to_check=good_consensus_bootstrapped_rids - + bad_consensus_bootstrapped_rids, + pp_good_rids=good_consensus_bootstrapped_rids, + pp_bad_rids=bad_consensus_bootstrapped_rids, ) ) self.waiting_queue.extend(good_reqs) @@ -1417,8 +1417,8 @@ def process_prealloc_queue(self: Scheduler, prealloc_rids: Optional[List[str]]): bad_consensus_prealloc_rids, ) = prealloc_rids good_reqs, failed_reqs = self.disagg_decode_prealloc_queue.pop_preallocated( - rids_to_check=good_consensus_prealloc_rids - + bad_consensus_prealloc_rids, + pp_good_rids=good_consensus_prealloc_rids, + pp_bad_rids=bad_consensus_prealloc_rids, ) self.disagg_decode_transfer_queue.extend(good_reqs) return [ diff --git a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py index e136befc4d54..0a904792b5af 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py +++ b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py @@ -40,6 +40,7 @@ def test_prealloc_abort_clears_receiver_before_removing_request(self): decode_req = SimpleNamespace(req=req, kv_receiver=receiver) queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue.pp_size = 1 queue.queue = [decode_req] queue.pending_reqs = [] queue.retracted_queue = [] @@ -86,6 +87,7 @@ def __eq__(self, other): decode_req = SimpleNamespace(req=req, kv_receiver=receiver) queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue.pp_size = 1 queue.queue = [decode_req] queue.pending_reqs = [decode_req] # same object, dual ownership queue.retracted_queue = [] diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index a5f052c7311e..762cb5489ccb 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -119,6 +119,7 @@ def _new_decode_req(self, rid: str, priority: int, *, failed: bool = False): def _new_queue(self, decode_reqs, *, low_priority_values_first: bool = False): queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue.pp_size = 1 queue.queue = list(decode_reqs) queue.pending_reqs = [] queue.retracted_queue = [] diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index 7cba0af202b7..4abdfdc1b933 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -294,6 +294,7 @@ def test_full_transfer_failure(self): def test_pop_preallocated_rechecks_budget_after_lock(self): queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue.pp_size = 1 req = MagicMock() req.rid = "req-1"