From cce5fe769638c6aeb95f3c3d50f8dcdf0223653f Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 18 Jul 2026 20:38:34 -0700 Subject: [PATCH] [Scheduler] Move the WAR barrier to right after each `run_batch` launch (#31687) --- python/sglang/srt/disaggregation/decode.py | 3 +-- python/sglang/srt/disaggregation/prefill.py | 3 +-- python/sglang/srt/managers/scheduler.py | 14 +++++++------- 3 files changed, 9 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index ec4f47707..5224920c4 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -2034,8 +2034,6 @@ class SchedulerDisaggregationDecodeMixin: continue self.process_decode_queue() - self._apply_war_barrier() - # Get the next batch to run plan = self.get_next_disagg_decode_batch_to_run( running_batch=self.running_batch @@ -2057,6 +2055,7 @@ class SchedulerDisaggregationDecodeMixin: # Launch the current batch if batch: batch_result = self.run_batch(batch) + self._apply_war_barrier() self.result_queue.append((batch.copy(), batch_result)) else: batch_result = None diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 6e8bc0db0..1f4c947fa 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -554,8 +554,6 @@ class SchedulerDisaggregationPrefillMixin: self.disagg_prefill_bootstrap_queue.pop_bootstrapped() ) - self._apply_war_barrier() - # Get the next batch to run plan = self.get_next_disagg_prefill_batch_to_run( running_batch=self.running_batch, last_batch=self.last_batch @@ -572,6 +570,7 @@ class SchedulerDisaggregationPrefillMixin: if self.enable_staging: self.maybe_prefetch_staging_for_batch(batch) batch_result = self.run_batch(batch) + self._apply_war_barrier() self.result_queue.append((batch.copy(), batch_result)) else: batch_result = None diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 3b14bbcdc..b2305a99b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1483,11 +1483,11 @@ class Scheduler( dispatch_event_loop(self) def _apply_war_barrier(self): - # Wait for the prev forward to finish reading the shared buffers this - # iter's schedule will overwrite. Fast path: wait on the read-done event - # the forward published after its snapshot (non-spec: decode graph; - # spec: draft_extend), then clear it. Else fall back to whole-forward - # wait_stream. + # Called right after each launch: order later schedule_stream work + # (result processing, next iteration's writes) behind the forward's + # shared-buffer reads. Fast path: wait on the read-done event the + # forward published after its snapshot (non-spec: decode graph; spec: + # draft_extend), then clear it. Else whole-forward wait_stream. if not self._war_barrier_enabled: return runner = self.model_worker.war_fastpath_runner @@ -1554,8 +1554,6 @@ class Scheduler( if self._engine_paused: continue - self._apply_war_barrier() - # Get the next batch to run plan = self.get_next_batch_to_run( running_batch=self.running_batch, last_batch=self.last_batch @@ -1583,6 +1581,8 @@ class Scheduler( # Launch the current batch if batch: batch_result = self.run_batch(batch) + # Fence result processing behind this forward's shared reads. + self._apply_war_barrier() self.result_queue.append((batch.copy(), batch_result)) else: batch_result = None