From 5601b7139d5aebc974919bca1a31be208b761a75 Mon Sep 17 00:00:00 2001 From: Bingxu Chen Date: Fri, 29 May 2026 16:02:31 +0800 Subject: [PATCH] [core] Make overlap-schedule WAR barrier CUDA-only (#26646) --- python/sglang/srt/disaggregation/decode.py | 3 ++- python/sglang/srt/disaggregation/prefill.py | 3 ++- python/sglang/srt/managers/overlap_utils.py | 6 ++++-- python/sglang/srt/managers/scheduler.py | 6 +++++- 4 files changed, 13 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 5138f571a..b8086153f 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1629,7 +1629,8 @@ class SchedulerDisaggregationDecodeMixin: continue # WAR barrier: this iter's schedule writes to shared GPU buffers wait for prev forward's reads. - self.schedule_stream.wait_stream(self.forward_stream) + if self._war_barrier_enabled: + self.schedule_stream.wait_stream(self.forward_stream) # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index ca2237861..350c84cb4 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -437,7 +437,8 @@ class SchedulerDisaggregationPrefillMixin: continue # WAR barrier on shared GPU buffers (req_to_token_pool / SWA mapping). - self.schedule_stream.wait_stream(self.forward_stream) + if self._war_barrier_enabled: + self.schedule_stream.wait_stream(self.forward_stream) # Get the next batch to run batch = self.get_next_disagg_prefill_batch_to_run() diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 2c9b926d9..5f622d4ff 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -106,8 +106,10 @@ class FutureMap: (self.req_pool_size,), dtype=torch.int64, device=self.device ) # Pinned host copy of new_seq_lens_buf + private stream for fwd-prepare - # D2H pulls (gated only on publish, off the schedule stream). - if _is_cuda or _is_hip: + # D2H pulls (gated only on publish, off the schedule stream). CUDA-only: + # recovers occupancy lost to the WAR barrier (also CUDA-only); other + # platforms have no barrier and use the plain .cpu() bootstrap path. + if _is_cuda: self.new_seq_lens_cpu_pinned = torch.empty( (self.req_pool_size,), dtype=torch.int64, pin_memory=True ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ff6ea128e..723bd2d00 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -245,6 +245,7 @@ from sglang.srt.utils import ( get_available_gpu_memory, get_bool_env_var, get_int_env_var, + is_cuda, is_mps, kill_itself_when_parent_died, require_mlp_sync, @@ -1362,6 +1363,8 @@ class Scheduler( self.schedule_stream = self.device_module.Stream(priority=0) if self.device == "cpu": self.schedule_stream.synchronize = lambda: None # No-op for CPU + # WAR barrier is CUDA-only; other platforms keep the pre-barrier behavior. + self._war_barrier_enabled = is_cuda() with self.device_module.StreamContext(self.schedule_stream): dispatch_event_loop(self) @@ -1412,7 +1415,8 @@ class Scheduler( continue # WAR barrier: this iter's schedule writes to shared GPU buffers wait for prev forward's reads. - self.schedule_stream.wait_stream(self.forward_stream) + if self._war_barrier_enabled: + self.schedule_stream.wait_stream(self.forward_stream) # Get the next batch to run batch = self.get_next_batch_to_run()