diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 45c3c8c57..c339b88f0 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -4106,33 +4106,49 @@ class Scheduler( if not self.running_batch.is_empty(): self.running_batch.filter_batch() - if len(self.running_batch.reqs) != 0: - # Decode-side retract always rebootstraps (recomputes the KV from - # the prefill), so skip the device->host KV offload that release_req - # would otherwise do; the offloaded copy would be immediately - # discarded. Non-decode modes ignore offload_kv (they never offload). - retracted_reqs = self.running_batch.reqs - retract_all( - reqs=retracted_reqs, - server_args=self.server_args, - req_to_token_pool=self.running_batch.req_to_token_pool, - token_to_kv_pool_allocator=self.running_batch.token_to_kv_pool_allocator, - tree_cache=self.running_batch.tree_cache, - hisparse_coordinator=self.running_batch.hisparse_coordinator, - offload_kv=False, - ) - self.running_batch.reqs = [] - for req in retracted_reqs: - if self.disaggregation_mode == DisaggregationMode.DECODE: - if req.output_ids: - req.pd_rebootstrap_forced_output_id = req.output_ids.pop() - req.pd_rebootstrap_in_progress = True - req.time_stats.set_retract_time() - self.disagg_decode_prealloc_queue.hold_rebootstrap(req) - else: - self._add_request_to_queue(req) - self.running_batch.batch_is_full = False + retracted_reqs = list(self.running_batch.reqs) + if ( + self.chunked_req is not None + and not self.chunked_req.finished() + and self.chunked_req not in retracted_reqs + and self.disaggregation_mode != DisaggregationMode.PREFILL + ): + retracted_reqs.append(self.chunked_req) + + if retracted_reqs: + # Decode-side retract always rebootstraps (recomputes the KV from + # the prefill), so skip the device->host KV offload that release_req + # would otherwise do; the offloaded copy would be immediately + # discarded. Non-decode modes ignore offload_kv (they never offload). + retract_all( + reqs=retracted_reqs, + server_args=self.server_args, + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + tree_cache=self.tree_cache, + hisparse_coordinator=self.hisparse_coordinator, + offload_kv=False, + ) + self.running_batch.reqs = [] + for req in retracted_reqs: + if self.disaggregation_mode == DisaggregationMode.DECODE: + if req.output_ids: + req.pd_rebootstrap_forced_output_id = req.output_ids.pop() + req.pd_rebootstrap_in_progress = True + req.time_stats.set_retract_time() + self.disagg_decode_prealloc_queue.hold_rebootstrap(req) + else: + self._add_request_to_queue(req) + + self.running_batch.batch_is_full = False + # In disagg-PREFILL, keep a live mid-chunk chunked_req rather than retract it: + # freeing its KV under a live disagg KV-sender crashes pop_bootstrapped or + # sends freed/reused KV to decode. Kept, it resumes prefill after the pause. + # TODO(disagg-prefill-retract): tear the sender down (abort + release metadata + # buffer + reset pending_bootstrap) before freeing KV, then retract for real. + # Until then a weight-update pause leaves stale-weight prefix KV (off-policy). + if self.disaggregation_mode != DisaggregationMode.PREFILL: self.chunked_req = None # Surface the paused state to dashboards immediately. The scheduler diff --git a/test/registered/unit/managers/test_scheduler_pause_generation.py b/test/registered/unit/managers/test_scheduler_pause_generation.py index eac3c1feb..63883e07b 100644 --- a/test/registered/unit/managers/test_scheduler_pause_generation.py +++ b/test/registered/unit/managers/test_scheduler_pause_generation.py @@ -35,6 +35,7 @@ class TestSchedulerPauseGeneration(unittest.TestCase): scheduler.tree_cache = MagicMock() scheduler.tree_cache.protected_size.return_value = 0 scheduler.req_to_token_pool = MagicMock() + scheduler.hisparse_coordinator = MagicMock() scheduler.result_queue = deque() scheduler.disaggregation_mode = DisaggregationMode.NULL # Support _kv_snap diagnostic logging in patched schedulers @@ -153,6 +154,31 @@ class TestSchedulerPauseGeneration(unittest.TestCase): ) self.assertIsNone(scheduler.chunked_req) + def test_retract_fold_in_releases_via_scheduler_hisparse_coordinator(self): + """retract of a folded-in last extend batch must release through the scheduler-owned hisparse coordinator.""" + scheduler = self._new_scheduler() + scheduler.disaggregation_mode = DisaggregationMode.NULL + scheduler.waiting_queue = [] + scheduler._add_request_to_queue = MagicMock() + scheduler.server_args = MagicMock() + + req = MagicMock() + req.finished.return_value = False + req.req_pool_idx = None + last_batch = MagicMock() + last_batch.forward_mode.is_extend.return_value = True + last_batch.is_empty.return_value = False + last_batch.reqs = [req] + scheduler.last_batch = last_batch + + scheduler.pause_generation(PauseGenerationReqInput(mode="retract")) + + scheduler.hisparse_coordinator.retract_req.assert_called_once_with(req) + self.assertEqual( + [call.args[0] for call in scheduler._add_request_to_queue.call_args_list], + [req], + ) + def test_retract_empty_running_batch_requeues_nothing(self): """retract with empty running_batch must not release or requeue any request.""" scheduler = self._new_scheduler() @@ -165,6 +191,36 @@ class TestSchedulerPauseGeneration(unittest.TestCase): self.assertEqual(len(scheduler.waiting_queue), 0) self.assertIs(scheduler.running_batch.reqs, original_reqs) + def test_retract_empty_clears_chunked_req_and_batch_is_full(self): + """retract with everything empty must still clear chunked_req and batch_is_full.""" + scheduler = self._new_scheduler() + scheduler.waiting_queue = [] + scheduler.chunked_req = MagicMock() + scheduler.running_batch.batch_is_full = True + + scheduler.pause_generation(PauseGenerationReqInput(mode="retract")) + + self.assertIsNone(scheduler.chunked_req) + self.assertFalse(scheduler.running_batch.batch_is_full) + + def test_retract_disagg_prefill_keeps_live_chunked_req(self): + """disagg-PREFILL retract must leave a live mid-chunk chunked_req untouched.""" + scheduler = self._new_scheduler() + scheduler.disaggregation_mode = DisaggregationMode.PREFILL + scheduler._add_request_to_queue = MagicMock() + scheduler.last_batch = None + + chunked_req = MagicMock() + chunked_req.finished.return_value = False + scheduler.chunked_req = chunked_req + + with patch("sglang.srt.managers.scheduler.retract_all") as mock_retract_all: + scheduler.pause_generation(PauseGenerationReqInput(mode="retract")) + + mock_retract_all.assert_not_called() + scheduler._add_request_to_queue.assert_not_called() + self.assertIs(scheduler.chunked_req, chunked_req) + def test_retract_drains_overlap_queue(self): """retract with overlap enabled should drain the result_queue.""" scheduler = self._new_scheduler()