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