diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index ec0579ec1e24..524f0bdb6fea 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -898,7 +898,7 @@ def pop_preallocated( self._resolve_pending_reqs() 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 + rids_to_check = set(pp_good_rids) | set(pp_bad_rids) failed_reqs = [] preallocated_reqs = [] diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 862ba08f5ad3..861c13fb4e49 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -4224,6 +4224,8 @@ def abort_request(self, recv_req: AbortReq): if hasattr(req.disagg_kv_sender, "abort"): req.disagg_kv_sender.abort() + if self.ps.pp_size > 1: + prepare_abort(req, "Aborted by AbortReq.") # Abort in-flight requests for req in self.disagg_prefill_inflight_queue: @@ -4238,6 +4240,8 @@ def abort_request(self, recv_req: AbortReq): if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}") decode_req.kv_receiver.abort() + if self.ps.pp_size > 1: + prepare_abort(decode_req.req, "Aborted by AbortReq.") # Abort requests waiting for kvcache to release tree cache for decode_req in self.disagg_decode_transfer_queue.queue: diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index f3e1f76f4ad1..4e475f8a76a5 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -24,7 +24,7 @@ set_is_extend_in_batch, ) from sglang.srt.managers.overlap_utils import RelayPayload -from sglang.srt.managers.schedule_batch import Req, ScheduleBatch +from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req, ScheduleBatch from sglang.srt.managers.utils import ( GenerationBatchResult, get_logprob_dict_from_result, @@ -847,6 +847,18 @@ def _pp_pd_get_bootstrapped_ids(self: Scheduler): bad_bootstrapped_rids = list( set(prev_bad_bootstrapped_rids) | set(curr_bad_bootstrapped_rids) ) + # Route locally-aborted reqs through the bad-union consensus so every PP + # rank flushes them in the same consensus round, regardless of when the + # AbortReq reaches each rank and regardless of whether + # disagg_kv_sender.abort() drives the poll to Failed (it is optional). + aborted_rids = { + req.rid + for req in self.disagg_prefill_bootstrap_queue.queue + if isinstance(req.finished_reason, FINISH_ABORT) + } + good_bootstrapped_rids, bad_bootstrapped_rids = self._route_aborts_to_bad( + good_bootstrapped_rids, bad_bootstrapped_rids, aborted_rids + ) return [good_bootstrapped_rids, bad_bootstrapped_rids] def _pp_pd_get_prefill_transferred_ids(self: Scheduler): @@ -1369,8 +1381,31 @@ def _pp_pd_get_prealloc_ids(self: Scheduler): bad_prealloc_rids = list( set(prev_bad_prealloc_rids) | set(curr_bad_prealloc_rids) ) + # Same abort routing as the prefill bootstrap consensus above. + aborted_rids = { + decode_req.req.rid + for decode_req in self.disagg_decode_prealloc_queue.queue + if isinstance(decode_req.req.finished_reason, FINISH_ABORT) + } + good_prealloc_rids, bad_prealloc_rids = self._route_aborts_to_bad( + good_prealloc_rids, bad_prealloc_rids, aborted_rids + ) return [good_prealloc_rids, bad_prealloc_rids] + @staticmethod + def _route_aborts_to_bad(good_rids, bad_rids, aborted_rids): + """Move aborted rids out of the good (intersection) set and into the + bad (union) set, so PP consensus fails them uniformly on every rank. + + This also flushes aborted reqs that never reached good/bad consensus + (e.g. stuck in Bootstrapping with a sender that has no working abort()). + """ + if not aborted_rids: + return good_rids, bad_rids + good_rids = [rid for rid in good_rids if rid not in aborted_rids] + bad_rids = list(set(bad_rids) | set(aborted_rids)) + return good_rids, bad_rids + def _pp_pd_get_decode_transferred_ids(self: Scheduler): # get the current stage transfer success if self.pp_group.is_first_rank: