From 0cae88e7fb29fcb853075f236c4ca94721d8e798 Mon Sep 17 00:00:00 2001 From: Sohom Chakraborty <16609933+sohom-cs@users.noreply.github.com> Date: Wed, 30 Sep 2026 18:12:37 -0400 Subject: [PATCH] [Fix] Publish KV events while the engine is paused The paused branch of every scheduler event loop skips on_idle, so events recorded while paused were only published after continue_generation. In the RL weight-update sequence (pause in retract mode, update weights with flush_cache, continue), the AllBlocksCleared from the flush therefore reached KV-event routers only after resume, and they kept routing against the old tree for the whole update. Publish from _record_scheduler_state_for_paused_engine, which every pausable loop calls on each paused iteration. Signed-off-by: Sohom Chakraborty <16609933+sohom-cs@users.noreply.github.com> --- python/sglang/srt/managers/scheduler.py | 3 ++ .../test_scheduler_pause_generation.py | 41 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 86947616f234..af1732bb344a 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -5299,6 +5299,9 @@ def _record_scheduler_state_for_paused_engine(self) -> None: self.metrics_reporter.record_scheduler_idle() else: self.metrics_reporter.record_scheduler_active(time.monotonic_ns()) + # Paused loops skip on_idle; publish here so a pause-time flush_cache + # (weight update) reaches KV-event routers before resume. + self.kv_events_publisher.publish_kv_events() def is_fully_idle(self, for_health_check=False, ignore_waiting=False) -> bool: # Health check piggybacks on running requests in process_output. diff --git a/test/registered/unit/managers/test_scheduler_pause_generation.py b/test/registered/unit/managers/test_scheduler_pause_generation.py index 2e2faf752854..f1950758ffc9 100644 --- a/test/registered/unit/managers/test_scheduler_pause_generation.py +++ b/test/registered/unit/managers/test_scheduler_pause_generation.py @@ -24,7 +24,11 @@ ) from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.scheduler_components.kv_events_publisher import ( + SchedulerKvEventsPublisher, +) from sglang.srt.managers.scheduler_components.pool_stats_observer import PoolStats +from sglang.srt.mem_cache.radix_cache import RadixCache from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.runtime_context import publish, reset_context from sglang.srt.sampling.sampling_params import SamplingParams @@ -294,6 +298,43 @@ def test_inplace_only_sets_flag(self): self.assertIs(scheduler.cur_batch_for_debug, original_cur_batch) self.assertIs(scheduler.chunked_req, original_chunked_req) + def test_cache_flush_while_paused_is_published_before_resume(self): + """A cache flush during a retract pause (the weight-update window) must be + published by the paused loop, not held until continue_generation.""" + scheduler = self._new_scheduler() + scheduler.tree_cache = RadixCache.create_simulated(enable_kv_cache_events=True) + scheduler.tree_cache.take_events() # Drop the construction-time clear. + published = [] + scheduler.kv_events_publisher = SchedulerKvEventsPublisher( + kv_events_config=None, + tree_cache=scheduler.tree_cache, + send_metrics_from_scheduler=None, + max_running_requests=1, + max_total_num_tokens=1, + get_stats=lambda: None, + ) + scheduler.kv_events_publisher.enable_kv_cache_events = True + scheduler.kv_events_publisher.kv_event_publisher = SimpleNamespace( + publish=published.append + ) + scheduler.grammar_manager = MagicMock(grammar_queue=[]) + scheduler.dllm_manager = MagicMock() + scheduler.dllm_manager.any_staging_reqs.return_value = False + scheduler.enable_hisparse = False + scheduler.enable_hierarchical_cache = False + scheduler.enable_lmcache = False + scheduler.draft_worker = None + + scheduler.pause_generation(PauseGenerationReqInput(mode="retract")) + self.assertTrue(scheduler.flush_cache(empty_cache=False)) + scheduler._record_scheduler_state_for_paused_engine() + + self.assertEqual( + [type(e).__name__ for batch in published for e in batch.events], + ["AllBlocksCleared"], + "flush_cache while paused was not published by the paused loop", + ) + def test_inplace_does_not_drain_overlap_queue(self): """in_place should not process the overlap result_queue.""" scheduler = self._new_scheduler()