diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 5c8dbd76b91a..f22b02adc77c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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. 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 1dff4fde4777..9ec25b53de38 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/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 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 f2463febaea5..9e20b3004eb6 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 @@ -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): 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 fe81e4c4a3df..c09305ce4b71 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 @@ -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()