Fix missed hisparse release and stale field cleanup in pause retract (#30674)

2026-07-15 14:31:11 +08:00
parent b6cc897fea
commit 1967b9ec99
2 changed files with 98 additions and 26 deletions
@@ -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
@@ -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()