From ff356ebae7696a27057b6834e6d4693e516580c4 Mon Sep 17 00:00:00 2001 From: Jin Pan Date: Wed, 23 Sep 2026 10:49:35 +0000 Subject: [PATCH] perf(amd): keep speculative overlap scheduling on the forward stream With overlap scheduling and speculative decoding on ROCm, the FutureMap publish waits block the host (publish_ready.synchronize()): an Event.wait() on the separate schedule stream leaves a cross-queue wait pending while the forward graph runs, and that slows every dispatch of the graph. Schedule on the forward stream instead, so stream order carries the forward-to-schedule dependencies and the publish waits stay in one queue. Without speculation there is no publish wait and the separate schedule stream keeps overlapping the next batch's preparation, so it stays. --- python/sglang/srt/managers/overlap_utils.py | 10 +++++++--- python/sglang/srt/managers/scheduler.py | 13 +++++++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index e43920cd412c..4593b456e559 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -266,6 +266,9 @@ def __init__( # full decision (per-backend flag + TBO / piecewise CG overrides). self.needs_cpu_seq_lens = needs_cpu_seq_lens self.needs_confidence_relay = needs_confidence_relay + # Set by the scheduler when it schedules on the forward stream (HIP with + # speculative decoding); publish waits then stay in that queue. + self.same_queue_publish = False self.req_pool_size = req_to_token_pool.req_to_token.shape[0] # Kept for the mixed-tail late binding (reserved-slot gather). self.req_to_token = req_to_token_pool.req_to_token @@ -475,7 +478,7 @@ def resolve_mixed_spec_tails(self, batch: ScheduleBatch) -> None: if n == 0: return if self.publish_ready is not None: - if _is_hip: + if _is_hip and not self.same_queue_publish: self.publish_ready.synchronize() else: self.publish_ready.wait() @@ -527,8 +530,9 @@ def resolve_seq_lens_cpu(self, batch: ScheduleBatch) -> None: # forward publish; a stale consume means a publish went missing. assert self._publish_fresh, "resolve without a fresh forward publish" self._publish_fresh = False - if _is_hip: - # Temporary workaround: Event.wait() regresses TPOT on AMD MI355. + if _is_hip and not self.same_queue_publish: + # A device wait left pending in another HIP queue while the forward + # graph runs slows every dispatch of that graph, so block the host. self.publish_ready.synchronize() else: self.publish_ready.wait() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f0d8533e03ff..c4b23aaaa04b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1863,6 +1863,19 @@ def run_event_loop(self) -> None: self.schedule_stream = allocate_distinct_stream( self.device_module, (self.forward_stream,) ) + if ( + _is_hip + and self.enable_overlap + and self.spec_algorithm.is_some() + and get_parallel().pp_size == 1 + ): + # A device wait pending in another HW queue slows every dispatch of the + # forward graph running meanwhile, so schedule on the forward queue and + # let stream order carry the forward-to-schedule dependencies instead. + # Without speculation there is no publish wait to remove, and the + # separate schedule stream still overlaps the next batch's preparation. + self.schedule_stream = self.forward_stream + self.future_map.same_queue_publish = True # The global WAR barrier fences the scheduler's next shared-buffer write # on the previous forward's read of the unified memory pool. self._war_barrier_enabled = is_cuda() or envs.SGLANG_ENABLE_WAR_BARRIER.get()