[Scheduler] Move the WAR barrier to right after each run_batch launch (#31687)

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

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

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
+        # Called right after each launch: order later schedule_stream work
-        # iter's schedule will overwrite. Fast path: wait on the read-done event
+        # (result processing, next iteration's writes) behind the forward's
-        # the forward published after its snapshot (non-spec: decode graph;
+        # shared-buffer reads. Fast path: wait on the read-done event the
-        # spec: draft_extend), then clear it. Else fall back to whole-forward
+        # forward published after its snapshot (non-spec: decode graph; spec:
-        # wait_stream.
+        # 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