[PD] Retract transferred decode requests on a retract pause

python/sglang/srt/managers/scheduler.py

@@ -5156,6 +5156,10 @@ def is_fully_idle(self, for_health_check=False, ignore_waiting=False) -> bool:
         # Waiting queues: waiting + bootstrapping + preallocation + kv transfer (decode)
         if not ignore_waiting:
             idle &= len(self.waiting_queue) == 0
+        else:
+            # A PD-decode request waiting with its transferred KV still owns
+            # its request row; a flush would free it.
+            idle &= not any(r.kv.holds_kv for r in self.waiting_queue)
 
         if (
             for_health_check
@@ -5618,6 +5622,16 @@ def pause_generation(self, recv_req: PauseGenerationReqInput):
         ):
             retract_reqs.append(self.chunked_req)
 
+        if self.disaggregation_mode == DisaggregationMode.DECODE:
+            # Transferred requests wait for a slot holding their row and KV;
+            # retract them too, so their KV is rebuilt under the new weights.
+            staged = [r for r in self.waiting_queue if r.kv.holds_kv]
+            if staged:
+                self.waiting_queue = [
+                    r for r in self.waiting_queue if not r.kv.holds_kv
+                ]
+                retract_reqs += staged
+
         self.last_batch = None
         self.cur_batch_for_debug = None
 

test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py

@@ -13,9 +13,11 @@
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
 
 
-def _make_req(is_retracted: bool) -> MagicMock:
+def _make_req(is_retracted: bool, holds_kv: bool = False) -> MagicMock:
     req = MagicMock()
     req.is_retracted = is_retracted
+    # Retracted and never-scheduled requests hold no request row.
+    req.kv.holds_kv = holds_kv
     return req
 
 
@@ -85,6 +87,18 @@ def test_flush_cache_succeeds_on_mixed_queue_while_paused(self):
         self.assertTrue(scheduler.flush_cache())
         scheduler.tree_cache.reset.assert_called_once()
 
+    def test_paused_flush_refuses_waiting_req_holding_kv(self):
+        """A PD-decode request whose KV has arrived waits holding its request
+        row; a paused flush would free KV it still owns."""
+        scheduler = self._new_scheduler()
+        scheduler.waiting_queue = [
+            _make_req(is_retracted=False),
+            _make_req(is_retracted=False, holds_kv=True),
+        ]
+
+        self.assertFalse(scheduler.flush_cache())
+        scheduler.tree_cache.reset.assert_not_called()
+
     def test_is_fully_idle_default_ignores_nothing(self):
         """Regression guard: every other is_fully_idle() caller passes no args,
         so a non-empty waiting_queue must still report not-idle by default."""

test/registered/unit/managers/test_scheduler_pause_generation.py

@@ -647,6 +647,39 @@ def test_pd_decode_retract_requeues_for_rebootstrap(self):
         mock_retract_all.assert_called_once()
         self.assertEqual(mock_retract_all.call_args.kwargs["offload_kv"], False)
 
+    def test_pd_decode_retract_includes_transferred_waiting_reqs(self):
+        """A transferred request waiting for a running slot holds a request row
+        and KV; retract must rebuild it too, so a paused flush cannot free it."""
+        scheduler = self._new_scheduler()
+        scheduler.disaggregation_mode = DisaggregationMode.DECODE
+        scheduler.disagg_decode_prealloc_queue = MagicMock()
+
+        def _req(rid, req_pool_idx):
+            return SimpleNamespace(
+                rid=rid,
+                finished=lambda: False,
+                output_ids=[7],
+                time_stats=MagicMock(),
+                kv=SimpleNamespace(holds_kv=req_pool_idx is not None),
+            )
+
+        running = _req("running", 0)
+        staged = _req("staged", 1)
+        unstaged = _req("unstaged", None)
+        scheduler.running_batch.reqs = [running]
+        scheduler.waiting_queue = [staged, unstaged]
+
+        with patch("sglang.srt.managers.scheduler.retract_all") as mock_retract_all:
+            scheduler.pause_generation(PauseGenerationReqInput(mode="retract"))
+
+        self.assertEqual(mock_retract_all.call_args.kwargs["reqs"], [running, staged])
+        self.assertEqual(scheduler.waiting_queue, [unstaged])
+        held = [
+            c.args[0]
+            for c in scheduler.disagg_decode_prealloc_queue.hold_rebootstrap.call_args_list
+        ]
+        self.assertEqual(held, [running, staged])
+
     def test_pd_decode_retract_completes_offload_before_free(self):
         """PD decode retract must not free KV an offload copy is still reading, and
         must not leave offload acks that keep the paused engine from going idle."""