From 6617b4bc6647c3e0dcdb145a180f87f748cd5329 Mon Sep 17 00:00:00 2001 From: jeremyzhang866 Date: Fri, 11 Sep 2026 11:53:38 +0800 Subject: [PATCH 1/3] Fix final-prefill abort commit race --- python/sglang/srt/managers/scheduler.py | 14 ++++++++ .../batch_result_processor.py | 33 ++++++++++--------- 2 files changed, 31 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f0914b18d553..68be0ac7ddb3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3561,6 +3561,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. diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 9fa4e9b85c15..bac24ad9ad56 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -304,19 +304,13 @@ def process_batch_result_prefill( logprob_pt = 0 for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)): - should_commit_output = ( - not req.finished() - and not req.is_retracted - and req.inflight_middle_chunks <= 0 - ) + 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: - assert logits_output is not None statuses = logits_output.next_token_sampling_mask_status status = None if statuses is None else statuses[i] - sampling_mask_finish_reason = self.get_sampling_mask_finish_reason( - status=status - ) + sampling_mask_finish_reason = self.get_sampling_mask_finish_reason(status=status) if ( batch.return_hidden_states and logits_output.hidden_states is not None @@ -329,7 +323,7 @@ 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 ), ) @@ -353,12 +347,15 @@ 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) - self._maybe_update_reasoning_tokens(req, next_token_id) - req.update_finish_state() # A mixed spec tail committed its pending bonus token; advance # so the next spec prepare_for_decode reserves from the right base. @@ -384,7 +381,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: @@ -396,16 +393,16 @@ 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, @@ -875,10 +872,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 From 2700bc0971c7efb18957780e496342e9140addc7 Mon Sep 17 00:00:00 2001 From: jeremyzhang866 Date: Fri, 11 Sep 2026 14:18:44 +0800 Subject: [PATCH 2/3] Fix final-prefill abort commit race --- .../batch_result_processor.py | 30 +++- ...st_batch_result_processor_hidden_states.py | 135 ++++++++++++++++-- .../test_scheduler_chunked_req_gate.py | 75 +++++++++- 3 files changed, 224 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index bac24ad9ad56..d30b45ba9d42 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -304,13 +304,26 @@ 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 + 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] - sampling_mask_finish_reason = self.get_sampling_mask_finish_reason(status=status) + sampling_mask_finish_reason = self.get_sampling_mask_finish_reason( + status=status + ) if ( batch.return_hidden_states and logits_output.hidden_states is not None @@ -323,7 +336,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 and not drop_prefill_result + should_commit_output + and sampling_mask_finish_reason is None + and not drop_prefill_result ), ) @@ -355,7 +370,9 @@ def process_batch_result_prefill( else: # req output_ids are set here req.output_ids.append(next_token_id) + self._maybe_update_reasoning_tokens(req, next_token_id) + req.update_finish_state() # A mixed spec tail committed its pending bonus token; advance # so the next spec prepare_for_decode reserves from the right base. @@ -393,7 +410,8 @@ 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 and not drop_prefill_result, + store=sampling_mask_finish_reason is None + and not drop_prefill_result, ) if sampling_mask_finish_reason is not None: diff --git a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py index 363515e10343..b3612af443ac 100644 --- a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py +++ b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py @@ -3,7 +3,6 @@ from unittest.mock import Mock, patch import torch - from sglang.srt.environ import envs from sglang.srt.layers.logits_processor import LogitsProcessorOutput, SamplingMaskStatus from sglang.srt.managers.scheduler_components.batch_result_processor import ( @@ -18,7 +17,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, @@ -44,7 +45,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, @@ -77,26 +78,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: @@ -192,6 +206,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): diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index d13655d90cd0..879fa103d854 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/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 @@ -6,7 +6,6 @@ from unittest.mock import MagicMock import torch - from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel @@ -107,6 +106,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 @@ -179,5 +179,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() From a66277d84b38fa1c4ae8e668fcdb61d2b9c9d6d3 Mon Sep 17 00:00:00 2001 From: jeremyzhang866 Date: Thu, 24 Sep 2026 13:44:43 +0800 Subject: [PATCH 3/3] Fix import spacing for lint --- .../unit/managers/test_batch_result_processor_hidden_states.py | 1 + test/registered/unit/managers/test_scheduler_chunked_req_gate.py | 1 + 2 files changed, 2 insertions(+) diff --git a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py index b3612af443ac..2113de2f7a88 100644 --- a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py +++ b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py @@ -3,6 +3,7 @@ from unittest.mock import Mock, patch import torch + from sglang.srt.environ import envs from sglang.srt.layers.logits_processor import LogitsProcessorOutput, SamplingMaskStatus from sglang.srt.managers.scheduler_components.batch_result_processor import ( diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index ba490e864150..8a6c63dc1d04 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock import torch + from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel