[Perf][ROCm] Dual-stream decode with hipgraphs (#52033)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
Simon Danielsson committed
Aug 30, 2026 at 10:24 UTC
2c7d7dd64a2eaba0feedf42cab2f527486d7479c
2 files changed
+70
-57
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+29
-18
@@ -283,11 +283,24 @@ class MoERunner(MoERunnerInterface):
283
self._shared_experts: SharedExperts | None = None
284
if shared_experts is not None:
285
can_overlap = lambda: self._quant_method.mk_can_overlap_shared_experts
286
+ # When unquantized, shared expert inputs alias the hidden states,
287
+ # which can lead to race condition in multi-stream mode. Quantized
288
+ # routed experts copy the input into a fresh buffer first, breaking
289
+ # the alias, so overlap is safe. Only observed on ROCm.
290
+ routed_input_is_quantized = lambda: (
291
+ self.routed_experts.quant_method.moe_quant_config is not None
292
+ and self.routed_experts.quant_method.moe_quant_config.quant_dtype
293
+ is not None
294
+ )
295
+ is_multistream_safe = lambda: (
296
+ not current_platform.is_rocm() or routed_input_is_quantized()
297
+ )
298
self._shared_experts = SharedExperts(
299
shared_experts,
300
moe_config=moe_config,
301
enable_dbo=enable_dbo,
302
mk_can_overlap_shared_experts=can_overlap,
303
+ is_multistream_safe=is_multistream_safe,
304
)
305
306
# Needed for string -> MoERunner layer lookup in custom ops.
@@ -582,12 +595,17 @@ class MoERunner(MoERunnerInterface):
595
router_logits: torch.Tensor,
596
shared_experts_input: torch.Tensor | None,
597
input_ids: torch.Tensor | None = None,
598
+ shared_experts_overlapping: bool = False,
599
) -> tuple[torch.Tensor | None, torch.Tensor | UnfinalizedMoEOutput]:
600
"""Run expert routing and the fused MoE kernel via the quant method.
601
602
Orchestrates shared expert execution (before/after), expert selection
603
via the router, and the actual fused MoE computation. Returns
604
(shared_expert_output, fused_expert_output).
605
+
606
+ `shared_experts_overlapping` should be True only if using multi-stream
607
+ overlap. Then the shared expert was already launched in a separate
608
+ stream, so the results only have to be awaited here.
609
"""
610
self._maybe_apply_shared_experts(
611
shared_experts_input, SharedExpertsOrder.NO_OVERLAP
@@ -617,10 +635,9 @@ class MoERunner(MoERunnerInterface):
635
shared_experts_input=shared_experts_input,
636
)
637
620
- self._maybe_apply_shared_experts(
621
- shared_experts_input,
622
- SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
623
- )
638
+ if shared_experts_overlapping:
639
+ assert self._shared_experts is not None
640
+ self._shared_experts.wait()
641
642
return (
643
self._shared_experts.output if self._shared_experts is not None else None,
@@ -642,18 +659,6 @@ class MoERunner(MoERunnerInterface):
659
else nullcontext()
660
)
661
645
- def _maybe_sync_shared_experts_stream(
646
- self,
647
- shared_experts_input: torch.Tensor | None,
648
- ):
649
- # If router/gate provided, then apply it here.
650
- # (Note: This code runs only when "overlapped mode" is on to allow
651
- # parallel execution of shared experts with the RoutedExperts via
652
- # separate cuda stream)
653
- if self._shared_experts is not None:
654
- assert shared_experts_input is not None
655
- self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
656
-
662
def _maybe_add_zero_expert_output(
663
self,
664
result: torch.Tensor,
@@ -878,8 +883,13 @@ class MoERunner(MoERunnerInterface):
883
# TODO(bnell): this can be removed after MK migration is complete.
884
self.routed_experts._ensure_moe_quant_config_init()
885
881
- # Sync aux and main stream for shared expert multi-stream overlap.
882
- self._maybe_sync_shared_experts_stream(shared_experts_input)
886
+ # If using multi-stream overlap for shared experts, we must launch it
887
+ # before routed expert dispatch.
888
+ shared_experts_overlapping = False
889
+ if self._shared_experts is not None:
890
+ shared_experts_overlapping = self._shared_experts.maybe_forward_async(
891
+ shared_experts_input
892
+ )
893
894
# If the Runner holds the gate, apply it after the stream sync,
895
# so it can run overlapped with the
@@ -905,6 +915,7 @@ class MoERunner(MoERunnerInterface):
915
router_logits=router_logits,
916
shared_experts_input=shared_experts_input,
917
input_ids=input_ids,
918
+ shared_experts_overlapping=shared_experts_overlapping,
919
)
920
921
return self._maybe_combine(
vllm/model_executor/layers/fused_moe/runner/shared_experts.py
+41
-39
@@ -43,6 +43,7 @@ class SharedExperts(torch.nn.Module):
43
moe_config: FusedMoEConfig,
44
enable_dbo: bool,
45
mk_can_overlap_shared_experts: Callable[[], bool],
46
+ is_multistream_safe: Callable[[], bool],
47
):
48
super().__init__()
49
@@ -57,6 +58,10 @@ class SharedExperts(torch.nn.Module):
58
59
self._mk_can_overlap_shared_experts = mk_can_overlap_shared_experts
60
61
+ # Might not be safe to run multi-stream mode if routed and shared experts
62
+ # alias the same inputs
63
+ self._is_multistream_safe = is_multistream_safe
64
+
65
# Allow disabling of the separate shared experts stream for
66
# debug purposes.
67
# TODO: Remove this after more extensive testings with TP/DP
@@ -65,11 +70,12 @@ class SharedExperts(torch.nn.Module):
70
logger.debug_once("Disabling MoE shared_experts cuda stream")
71
self._stream = None
72
else:
68
- # TODO(rob): enable shared expert overlap with non-cuda-alike.
69
- # aux_stream() returns None on non-cuda-alike platforms.
73
self._stream = aux_stream()
74
if self._stream is not None:
75
logger.debug_once("Enabled separate cuda stream for MoE shared_experts")
76
+ # One pair per DBO ubatch id to sync aux and main stream.
77
+ self._input_ready_event = [torch.cuda.Event(), torch.cuda.Event()]
78
+ self._output_ready_event = [torch.cuda.Event(), torch.cuda.Event()]
79
80
# TODO(bnell): Hack for elastic_ep. Get rid of this
81
def _set_moe_config(self, new_moe_config: FusedMoEConfig):
@@ -106,11 +112,19 @@ class SharedExperts(torch.nn.Module):
112
if self._mk_can_overlap_shared_experts():
113
return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED
114
115
+ # On ROCm, empirically only DP-only deployments benefit from the overlap.
116
+ overlap_is_beneficial = not current_platform.is_rocm() or (
117
+ self._moe_config.moe_parallel_config.dp_size > 1
118
+ and self._moe_config.moe_parallel_config.tp_size == 1
119
+ )
120
+
121
should_run_shared_in_aux_stream = (
110
- current_platform.is_cuda()
122
+ current_platform.is_cuda_alike()
123
and self._stream is not None
124
and hidden_states.shape[0]
125
<= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
126
+ and overlap_is_beneficial
127
+ and self._is_multistream_safe()
128
)
129
130
if should_run_shared_in_aux_stream:
@@ -118,38 +132,31 @@ class SharedExperts(torch.nn.Module):
132
else:
133
return SharedExpertsOrder.NO_OVERLAP
134
121
- def maybe_sync_shared_experts_stream(
122
- self,
123
- shared_experts_input: torch.Tensor,
124
- ):
125
- experts_order = self._determine_shared_experts_order(shared_experts_input)
126
-
127
- if experts_order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED:
128
- assert self._stream is not None
129
-
130
- # Record that the clone will be used by shared_experts_stream
131
- # to avoid gc issue from deallocation of hidden_states_clone
132
- # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501
133
- # NOTE: We don't need shared_output.record_stream(current_stream())
134
- # because we synch the streams before using shared_output.
135
- shared_experts_input.record_stream(self._stream)
136
-
137
- # Mark sync start point for the aux stream since we will
138
- # run in parallel with router/gate.
139
- self._stream.wait_stream(current_stream())
140
-
141
- def _run_in_aux_stream(
142
- self,
143
- shared_experts_input: torch.Tensor,
144
- ) -> torch.Tensor:
145
- # TODO: assert that maybe_sync_shared_experts_stream has been called.
146
-
147
- # Run shared experts in parallel on a separate stream.
135
+ def maybe_forward_async(self, shared_experts_input: torch.Tensor) -> bool:
136
+ """Enqueue shared experts on the aux stream without waiting for them.
137
+
138
+ Returns true if the shared experts were enqueued, false otherwise. Call
139
+ `wait` to wait for the shared experts to finish if this returns true.
140
+ """
141
+ if (
142
+ self._determine_shared_experts_order(shared_experts_input)
143
+ != SharedExpertsOrder.MULTI_STREAM_OVERLAPPED
144
+ ):
145
+ return False
146
+ assert self._stream is not None
147
+ idx = self._output_idx
148
+ assert self._output[idx] is None
149
+ self._input_ready_event[idx].record(current_stream())
150
with torch.cuda.stream(self._stream):
149
- output = self._layer(shared_experts_input)
150
- current_stream().wait_stream(self._stream)
151
+ self._input_ready_event[idx].wait(self._stream)
152
+ self._output[idx] = self._layer(shared_experts_input)
153
+ self._output_ready_event[idx].record(self._stream)
154
+ return True
155
152
- return output
156
+ def wait(self) -> None:
157
+ """Block the main stream until `maybe_forward_async` output is ready."""
158
+ assert self._stream is not None
159
+ self._output_ready_event[self._output_idx].wait(current_stream())
160
161
@property
162
def _output_idx(self) -> int:
@@ -174,11 +181,6 @@ class SharedExperts(torch.nn.Module):
181
182
assert self._output[self._output_idx] is None
183
177
- if order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED:
178
- self._output[self._output_idx] = self._run_in_aux_stream(
179
- shared_experts_input
180
- )
181
- else:
182
- self._output[self._output_idx] = self._layer(shared_experts_input)
184
+ self._output[self._output_idx] = self._layer(shared_experts_input)
185
186
assert self._output[self._output_idx] is not None