100 lines
4.8 KiB
Diff
100 lines
4.8 KiB
Diff
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:
|