From 0a190d1c977111aa7530a4ac1c4b33658356d66c Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 7 Jun 2026 17:38:26 +0800 Subject: [PATCH] Fix PP is_fully_idle missing in-flight microbatches (#27446) --- python/sglang/srt/managers/scheduler.py | 9 ++++- .../test/scripted_runtime/context/queries.py | 13 ++++--- .../test/scripted_runtime/scheduler_hook.py | 15 ++++---- .../test_scripted_core_4gpu.py | 35 +++++++++++++++++++ 4 files changed, 58 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index da66e6cff..96e21ec75 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/python/sglang/test/scripted_runtime/context/queries.py b/python/sglang/test/scripted_runtime/context/queries.py index 3b114f1cd..45632a731 100644 --- a/python/sglang/test/scripted_runtime/context/queries.py +++ b/python/sglang/test/scripted_runtime/context/queries.py @@ -12,10 +12,15 @@ 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.running_batch is not None: - yield from s.running_batch.reqs - if s.last_batch is not None: - yield from s.last_batch.reqs + 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: + yield from s.last_batch.reqs def list_active_reqs(ctx: "ScriptedContext") -> List["Req"]: diff --git a/python/sglang/test/scripted_runtime/scheduler_hook.py b/python/sglang/test/scripted_runtime/scheduler_hook.py index d8172060e..494638f30 100644 --- a/python/sglang/test/scripted_runtime/scheduler_hook.py +++ b/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 ( - scheduler.chunked_req is None - and len(scheduler.waiting_queue) == 0 - and scheduler.running_batch.is_empty() - ): + if scheduler.is_fully_idle(): break - - server_args = scheduler.server_args - for _ in range(2 * (server_args.pp_size + server_args.pp_async_batch_depth)): - yield + else: + raise RuntimeError( + "scripted_runtime reset: scheduler did not become fully idle " + f"within {RESET_DRAIN_MAX_STEPS} steps" + ) ctx.flush_cache() yield diff --git a/test/registered/chunked_prefill/test_scripted_core_4gpu.py b/test/registered/chunked_prefill/test_scripted_core_4gpu.py index 2ff60d749..c1e14f0c8 100644 --- a/test/registered/chunked_prefill/test_scripted_core_4gpu.py +++ b/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()