perf(amd): keep speculative overlap scheduling on the forward stream

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()

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()