[Scheduler] Fix final chunked-prefill abort commit race

python/sglang/srt/managers/scheduler.py

@@ -3693,6 +3693,20 @@ def get_next_batch_to_run(
             and last_batch
             and last_batch.forward_mode.is_extend()
         ):
+            if getattr(self, "result_queue", None):
+                # Abort can arrive while the final prefill result is still delayed.
+                # Keep that request out of optimistic decode; the result processor
+                # will consume the pending finish and release its resources.
+                chunked_req_to_exclude.update(
+                    req
+                    for req in last_batch.reqs
+                    if req.to_finish is not None
+                    and (
+                        not last_batch.decoding_reqs
+                        or req not in last_batch.decoding_reqs
+                    )
+                )
+
             if last_batch.chunked_req is not None:
                 # In the context pipeline parallelism, after the last chunk, the current microbatch still track outdated chunked_req.
                 # We need to discard it.

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -305,13 +305,20 @@ def process_batch_result_prefill(
             logprob_pt = 0
 
             for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
+                drop_prefill_result = req.to_finish is not None and (
+                    not batch.decoding_reqs or req not in batch.decoding_reqs
+                )
                 should_commit_output = (
                     not req.finished()
                     and not req.is_retracted
                     and req.inflight_middle_chunks <= 0
                 )
                 sampling_mask_finish_reason = None
-                if should_commit_output and req.return_sampling_mask:
+                if (
+                    should_commit_output
+                    and not drop_prefill_result
+                    and req.return_sampling_mask
+                ):
                     assert logits_output is not None
                     statuses = logits_output.next_token_sampling_mask_status
                     status = None if statuses is None else statuses[i]
@@ -330,7 +337,9 @@ def process_batch_result_prefill(
                         capture_hidden_mode=prefill_hidden_capture_mode,
                         extend_input_len=extend_input_len_per_req[i],
                         store=(
-                            should_commit_output and sampling_mask_finish_reason is None
+                            should_commit_output
+                            and sampling_mask_finish_reason is None
+                            and not drop_prefill_result
                         ),
                     )
 
@@ -354,6 +363,11 @@ def process_batch_result_prefill(
                         self.beam_coordinator.commit_prefill(
                             req, up_to_tick=batch.forward_iter
                         )
+                    elif drop_prefill_result:
+                        # Abort won the race with this delayed prefill result.
+                        # Finish and clean up without committing the sampled token
+                        # or any metadata derived from it.
+                        req.update_finish_state(new_accepted_len=0)
                     else:
                         # req output_ids are set here
                         req.output_ids.append(next_token_id)
@@ -385,7 +399,7 @@ def process_batch_result_prefill(
                         if get_memory().enable_hisparse:
                             self.hisparse_coordinator.admit_request_into_staging(req)
 
-                    if sampling_mask_finish_reason is None:
+                    if sampling_mask_finish_reason is None and not drop_prefill_result:
                         self._maybe_collect_customized_info(i, req, logits_output)
 
                     if batch.return_logprob:
@@ -397,16 +411,17 @@ def process_batch_result_prefill(
                             extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
                             next_token_ids=next_token_ids,
                             logprob_pt=logprob_pt,
-                            store=sampling_mask_finish_reason is None,
+                            store=sampling_mask_finish_reason is None
+                            and not drop_prefill_result,
                         )
 
                     if sampling_mask_finish_reason is not None:
                         continue
 
-                    if req.return_sampling_mask:
+                    if not drop_prefill_result and req.return_sampling_mask:
                         self.add_sampling_mask_return_values(i, req, logits_output)
 
-                    if req.grammar is not None:
+                    if not drop_prefill_result and req.grammar is not None:
                         self._apply_prefill_grammar(
                             req=req,
                             next_token_id=next_token_id,
@@ -877,10 +892,14 @@ def advance_grammar_fsm(
             # Extend: advance over the single token each completed-prefill req emitted
             # (mirrors process_batch_result_prefill's per-req token indexing).
             for i, req in enumerate(batch.reqs):
+                drop_prefill_result = req.to_finish is not None and (
+                    not batch.decoding_reqs or req not in batch.decoding_reqs
+                )
                 if (
                     req.grammar is None
                     or req.is_retracted
                     or req.finished()
+                    or drop_prefill_result
                     or req.inflight_middle_chunks > 0
                 ):
                     continue

test/registered/unit/managers/test_batch_result_processor_hidden_states.py

@@ -22,7 +22,9 @@
 register_cpu_ci(est_time=11, suite="base-a-test-cpu")
 
 
-def _make_processor(case, server_mode: str = "full") -> SchedulerBatchResultProcessor:
+def _make_processor(
+    case, server_mode: str = "full", *, logprob_result_processor=None
+) -> SchedulerBatchResultProcessor:
     # The server-side hidden-state ceiling is a bag leaf.
     override = get_context().override_server_args(
         enable_return_hidden_states=True,
@@ -48,7 +50,7 @@ def _make_processor(case, server_mode: str = "full") -> SchedulerBatchResultProc
         metrics_reporter=metrics_reporter,
         draft_worker=None,
         model_worker=Mock(),
-        logprob_result_processor=None,
+        logprob_result_processor=logprob_result_processor,
         output_streamer=Mock(),
         beam_coordinator=Mock(),
         abort_request=lambda *args, **kwargs: None,
@@ -155,26 +157,39 @@ def test_packed_ids_are_copied_before_per_request_slicing(self):
 
 
 class _PrefillReq:
-    def __init__(self, *, rid: str, inflight_middle_chunks: int, return_hidden_states):
+    def __init__(
+        self,
+        *,
+        rid: str,
+        inflight_middle_chunks: int,
+        return_hidden_states,
+        return_logprob: bool = False,
+        to_finish=None,
+    ):
         self.rid = rid
         self.inflight_middle_chunks = inflight_middle_chunks
         self.return_hidden_states = return_hidden_states
         self.hidden_states = []
         self.is_retracted = False
         self.output_ids = []
+        self.to_finish = to_finish
+        self.finished_reason = None
         self.time_stats = Mock()
-        self.return_logprob = False
+        self.return_logprob = return_logprob
         self.return_sampling_mask = False
         self.grammar = None
         self.require_reasoning = False
         self.customized_info = None
         self.beam_group = None
+        self.return_routed_experts = False
 
     def finished(self):
-        return False
+        return self.finished_reason is not None
 
-    def update_finish_state(self):
-        return None
+    def update_finish_state(self, new_accepted_len=1):
+        if self.to_finish is not None:
+            self.finished_reason = self.to_finish
+            self.to_finish = None
 
 
 class _DecodeReq:
@@ -270,6 +285,111 @@ def test_active_middle_chunk_advances_before_new_last_request(self):
                 self.assertEqual(middle.hidden_states, [])
                 self.assertEqual(last.hidden_states, [[22.0]])
 
+    def test_aborted_final_prefill_drops_metadata_but_advances_logprob_offset(self):
+        self._assert_aborted_final_prefill()
+
+    def test_pending_abort_wins_over_sampling_mask_overflow(self):
+        self._assert_aborted_final_prefill(SamplingMaskStatus.OVERFLOW)
+
+    def _assert_aborted_final_prefill(self, sampling_status=None):
+        aborted = _PrefillReq(
+            rid="aborted",
+            inflight_middle_chunks=0,
+            return_hidden_states="last",
+            return_logprob=True,
+            to_finish=object(),
+        )
+        abort_reason = aborted.to_finish
+        aborted.return_sampling_mask = sampling_status is not None
+        live = _PrefillReq(
+            rid="live",
+            inflight_middle_chunks=0,
+            return_hidden_states="last",
+            return_logprob=True,
+        )
+        batch = SimpleNamespace(
+            reqs=[aborted, live],
+            decoding_reqs=[],
+            return_logprob=True,
+            return_hidden_states=True,
+            return_hidden_states_mode=CaptureHiddenMode.LAST,
+            spec_info=None,
+            prefill_stats=None,
+            dp_cooperation_info=None,
+        )
+        logits_output = SimpleNamespace(
+            sampling_mask_output=None,
+            next_token_sampling_mask_status=[sampling_status, None],
+            hidden_states=torch.tensor([[10.0], [20.0]]),
+            customized_info={"tag": ["abort-meta", "live-meta"]},
+            next_token_logprobs=None,
+            input_token_logprobs=torch.tensor([-1.0, -2.0, -10.0, -11.0, -12.0]),
+            next_token_top_logprobs_val=[],
+            next_token_top_logprobs_idx=[],
+            next_token_token_ids_logprobs_val=[],
+        )
+        result = SimpleNamespace(
+            copy_done=None,
+            auxiliary_host_output=None,
+            routed_experts_output=None,
+            indexer_topk_output=None,
+            logits_output=logits_output,
+            next_token_ids=torch.tensor([101, 202]),
+            extend_input_len_per_req=[2, 3],
+            extend_logprob_start_len_per_req=[0, 0],
+            grammar_advanced=False,
+            can_run_cuda_graph=False,
+            skipped_output_comm=False,
+        )
+        logprob_processor = Mock()
+        logprob_processor.calculate_num_input_logprobs.side_effect = [2, 3]
+        processor = _make_processor(
+            self, "last", logprob_result_processor=logprob_processor
+        )
+
+        def assert_clean_before_stream(reqs, *_):
+            self.assertEqual(reqs[0].output_ids, [])
+            self.assertEqual(reqs[0].hidden_states, [])
+            self.assertIsNone(reqs[0].customized_info)
+
+        processor.output_streamer.stream_output.side_effect = assert_clean_before_stream
+
+        with (
+            patch(
+                "sglang.srt.managers.scheduler_components."
+                "batch_result_processor.release_kv_cache"
+            ),
+            patch(
+                "sglang.srt.managers.scheduler_components."
+                "batch_result_processor.maybe_cache_unfinished_req"
+            ),
+            patch(
+                "sglang.srt.managers.scheduler_components."
+                "batch_result_processor.get_memory",
+                return_value=SimpleNamespace(enable_hisparse=False),
+            ),
+        ):
+            processor.process_batch_result_prefill(batch, result)
+
+        self.assertTrue(aborted.finished())
+        self.assertIs(aborted.finished_reason, abort_reason)
+        self.assertEqual(aborted.output_ids, [])
+        self.assertEqual(aborted.hidden_states, [])
+        self.assertIsNone(aborted.customized_info)
+        self.assertEqual(live.output_ids, [202])
+        self.assertEqual(live.hidden_states, [[20.0]])
+        self.assertEqual(live.customized_info, {"tag": ["live-meta"]})
+
+        self.assertEqual(logprob_processor.calculate_num_input_logprobs.call_count, 2)
+        logprob_processor.add_logprob_return_values.assert_called_once_with(
+            1,
+            live,
+            2,
+            [101, 202],
+            3,
+            logits_output,
+        )
+
 
 class TestPrefillSkippedOutput(CustomTestCase):
     def test_sampling_mask_middle_chunk_does_not_require_logits_output(self):

test/registered/unit/managers/test_scheduler_chunked_req_gate.py

@@ -1,4 +1,4 @@
-"""Regression tests for the SWA chunked-req stash gate (#24252)."""
+"""Regression tests for scheduler prefill/decode transition gates."""
 
 import unittest
 from array import array
@@ -115,6 +115,7 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler:
     s.tree_cache = tree_cache
     s.chunked_req = chunked_req
     s._pending_chunked_abort_req = None
+    s.result_queue = []
     return s
 
 
@@ -187,5 +188,76 @@ def test_no_chunked_req_never_mutates_state(self):
         self.assertIsNone(s.chunked_req)
 
 
+class TestPendingPrefillAbortGate(CustomTestCase):
+    @staticmethod
+    def _make_last_batch():
+        aborted_prefill = MagicMock(rid="aborted-prefill", to_finish=object())
+        live_prefill = MagicMock(rid="live-prefill", to_finish=None)
+        aborted_decode = MagicMock(rid="aborted-decode", to_finish=object())
+        for req in (aborted_prefill, live_prefill, aborted_decode):
+            req.finished.return_value = False
+
+        last_batch = MagicMock()
+        last_batch.forward_mode.is_extend.return_value = True
+        last_batch.reqs = [aborted_prefill, live_prefill, aborted_decode]
+        last_batch.decoding_reqs = [aborted_decode]
+        last_batch.chunked_req = None
+        last_batch.is_prefill_only = False
+        last_batch.batch_size.side_effect = lambda: len(last_batch.reqs)
+        last_batch.is_empty.side_effect = lambda: not last_batch.reqs
+
+        def filter_batch(*, chunked_req_to_exclude):
+            excluded = set(chunked_req_to_exclude)
+            last_batch.reqs = [req for req in last_batch.reqs if req not in excluded]
+
+        last_batch.filter_batch.side_effect = filter_batch
+        return last_batch, aborted_prefill, live_prefill, aborted_decode
+
+    def _scheduled_reqs(self, *, has_pending_result):
+        scheduler = _scheduler_for_get_next_batch(
+            tree_cache=MagicMock(), chunked_req=None
+        )
+        scheduler.result_queue = [object()] if has_pending_result else []
+        last_batch, aborted_prefill, live_prefill, aborted_decode = (
+            self._make_last_batch()
+        )
+        scheduler.get_new_batch_prefill.side_effect = lambda running_batch: (
+            NextBatchPlan(batch_to_run=None, running_batch=running_batch)
+        )
+        scheduled_reqs = []
+
+        def update_running_batch(batch):
+            scheduled_reqs.extend(batch.reqs)
+            return batch
+
+        scheduler.update_running_batch.side_effect = update_running_batch
+
+        Scheduler.get_next_batch_to_run(
+            scheduler,
+            running_batch=scheduler.running_batch,
+            last_batch=last_batch,
+        )
+
+        return scheduled_reqs, aborted_prefill, live_prefill, aborted_decode
+
+    def test_abort_is_excluded_before_optimistic_decode(self):
+        scheduled_reqs, aborted_prefill, live_prefill, aborted_decode = (
+            self._scheduled_reqs(has_pending_result=True)
+        )
+
+        self.assertNotIn(aborted_prefill, scheduled_reqs)
+        self.assertIn(live_prefill, scheduled_reqs)
+        self.assertIn(aborted_decode, scheduled_reqs)
+
+    def test_abort_is_not_excluded_after_prefill_result_was_processed(self):
+        scheduled_reqs, aborted_prefill, live_prefill, aborted_decode = (
+            self._scheduled_reqs(has_pending_result=False)
+        )
+
+        self.assertIn(aborted_prefill, scheduled_reqs)
+        self.assertIn(live_prefill, scheduled_reqs)
+        self.assertIn(aborted_decode, scheduled_reqs)
+
+
 if __name__ == "__main__":
     unittest.main()