Fix PP is_fully_idle missing in-flight microbatches (#27446)

python/sglang/srt/managers/scheduler.py

@@ -3330,7 +3330,7 @@ class Scheduler(
             and (self.last_batch is None or self.last_batch.is_empty())
             and (self.cur_batch is None or self.cur_batch.is_empty())
             and (not self.enable_overlap or len(self.result_queue) == 0)
-            and (self.ps.pp_size == 1 or all(x.is_empty() for x in self.running_mbs))
+            and self._pp_microbatches_drained()
         )
 
         # Waiting queues: waiting + bootstrapping + preallocation + kv transfer (decode)
@@ -3366,6 +3366,13 @@ class Scheduler(
 
         return idle
 
+    def _pp_microbatches_drained(self) -> bool:
+        if self.ps.pp_size == 1:
+            return True
+        return all(x.is_empty() for x in self.running_mbs) and all(
+            mb is None or mb.is_empty() for mb in self.mbs
+        )
+
     def attach_hicache_storage_wrapped(
         self, recv_req: AttachHiCacheStorageReqInput
     ) -> AttachHiCacheStorageReqOutput:

python/sglang/test/scripted_runtime/context/queries.py

@@ -12,6 +12,11 @@ def _get_all_reqs(ctx: "ScriptedContext") -> Iterator["Req"]:
     if s.chunked_req is not None:
         yield s.chunked_req
     yield from s.waiting_queue
+    if s.ps.pp_size > 1:
+        for mb in (*s.mbs, *s.last_mbs, *s.running_mbs):
+            if mb is not None:
+                yield from mb.reqs
+    else:
         if s.running_batch is not None:
             yield from s.running_batch.reqs
         if s.last_batch is not None:

python/sglang/test/scripted_runtime/scheduler_hook.py

@@ -99,16 +99,13 @@ def _reset_engine_state(ctx: ScriptedContext) -> Generator:
     ctx.abort_all()
     for _ in range(RESET_DRAIN_MAX_STEPS):
         yield
-        if (
+        if scheduler.is_fully_idle():
-            scheduler.chunked_req is None
-            and len(scheduler.waiting_queue) == 0
-            and scheduler.running_batch.is_empty()
-        ):
             break
-
+    else:
-    server_args = scheduler.server_args
+        raise RuntimeError(
-    for _ in range(2 * (server_args.pp_size + server_args.pp_async_batch_depth)):
+            "scripted_runtime reset: scheduler did not become fully idle "
-        yield
+            f"within {RESET_DRAIN_MAX_STEPS} steps"
+        )
 
     ctx.flush_cache()
     yield

test/registered/chunked_prefill/test_scripted_core_4gpu.py

@@ -8,6 +8,7 @@ from sglang.test.scripted_runtime_chunked_helpers import (
     SMALL_MODEL,
     base_engine_kwargs,
     run_until_all_finished,
+    run_until_finished,
 )
 
 register_cuda_ci(est_time=900, stage="extra-b", runner_config="4-gpu-h100")
@@ -55,6 +56,40 @@ class TestScriptedPpChunkSweep(ScriptedTestCase):
                 f"req {r.rid!r} did not finish"
             )
 
+    def test_pp_flush_cache_during_inflight_chunk_results(self):
+        """flush_cache landing after the last chunk dispatch but before its batch result is processed must not corrupt the radix tree."""
+        self.server.execute_script(self._script_flush_during_inflight_chunk_results)
+
+    @staticmethod
+    def _script_flush_during_inflight_chunk_results(t: ScriptedContext):
+        scheduler = t.scheduler
+        r = t.start_req(prompt_len=2 * _CHUNK_SIZE - 3, max_new_tokens=2)
+
+        # Wait for the window where the queues and the current microbatch slot
+        # bindings are all clear, yet dispatched chunk batch results are still
+        # in flight in the PP pipeline, then post flush_cache into that window.
+        flushed = False
+        for _ in range(DEFAULT_MAX_STEPS):
+            in_flight = any(
+                mb is not None and not mb.is_empty() for mb in scheduler.mbs
+            )
+            queues_clear = (
+                scheduler.chunked_req is None
+                and len(scheduler.waiting_queue) == 0
+                and all(x.is_empty() for x in scheduler.running_mbs)
+                and (scheduler.cur_batch is None or scheduler.cur_batch.is_empty())
+                and (scheduler.last_batch is None or scheduler.last_batch.is_empty())
+            )
+            if in_flight and queues_clear:
+                t.flush_cache()
+                flushed = True
+                break
+            yield
+        assert flushed, "never observed in-flight chunk results with clear queues"
+
+        yield from run_until_finished(r, max_steps=DEFAULT_MAX_STEPS)
+        assert r.finished, f"req {r.rid!r} did not finish after flush_cache"
+
 
 if __name__ == "__main__":
     unittest.main()