[Perf][ROCm] Dual-stream decode with hipgraphs (#48223)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Simon Danielsson committed
Aug 12, 2026 at 03:32 UTC
47ececb58e073c6ff7a99033950455763f302055
2 files changed
+55
-55
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+16
-18
@@ -573,12 +573,17 @@ class MoERunner(MoERunnerInterface):
573
router_logits: torch.Tensor,
574
shared_experts_input: torch.Tensor | None,
575
input_ids: torch.Tensor | None = None,
576
+ shared_experts_overlapping: bool = False,
577
) -> tuple[torch.Tensor | None, torch.Tensor]:
578
"""Run expert routing and the fused MoE kernel via the quant method.
579
580
Orchestrates shared expert execution (before/after), expert selection
581
via the router, and the actual fused MoE computation. Returns
582
(shared_expert_output, fused_expert_output).
583
+
584
+ `shared_experts_overlapping` should be True only if using multi-stream
585
+ overlap. Then the shared expert was already launched in a separate
586
+ stream, so the results only have to be awaited here.
587
"""
588
self._maybe_apply_shared_experts(
589
shared_experts_input, SharedExpertsOrder.NO_OVERLAP
@@ -608,10 +613,9 @@ class MoERunner(MoERunnerInterface):
613
shared_experts_input=shared_experts_input,
614
)
615
611
- self._maybe_apply_shared_experts(
612
- shared_experts_input,
613
- SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
614
- )
616
+ if shared_experts_overlapping:
617
+ assert self._shared_experts is not None
618
+ self._shared_experts.wait()
619
620
return (
621
self._shared_experts.output if self._shared_experts is not None else None,
@@ -633,18 +637,6 @@ class MoERunner(MoERunnerInterface):
637
else nullcontext()
638
)
639
636
- def _maybe_sync_shared_experts_stream(
637
- self,
638
- shared_experts_input: torch.Tensor | None,
639
- ):
640
- # If router/gate provided, then apply it here.
641
- # (Note: This code runs only when "overlapped mode" is on to allow
642
- # parallel execution of shared experts with the RoutedExperts via
643
- # separate cuda stream)
644
- if self._shared_experts is not None:
645
- assert shared_experts_input is not None
646
- self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
647
-
640
def _maybe_add_zero_expert_output(
641
self,
642
result: torch.Tensor,
@@ -849,8 +841,13 @@ class MoERunner(MoERunnerInterface):
841
# TODO(bnell): this can be removed after MK migration is complete.
842
self.routed_experts._ensure_moe_quant_config_init()
843
852
- # Sync aux and main stream for shared expert multi-stream overlap.
853
- self._maybe_sync_shared_experts_stream(shared_experts_input)
844
+ # If using multi-stream overlap for shared experts, we must launch it
845
+ # before routed expert dispatch.
846
+ shared_experts_overlapping = False
847
+ if self._shared_experts is not None:
848
+ shared_experts_overlapping = self._shared_experts.maybe_forward_async(
849
+ shared_experts_input
850
+ )
851
852
# If the Runner holds the gate, apply it after the stream sync,
853
# so it can run overlapped with the
@@ -876,6 +873,7 @@ class MoERunner(MoERunnerInterface):
873
router_logits=router_logits,
874
shared_experts_input=shared_experts_input,
875
input_ids=input_ids,
876
+ shared_experts_overlapping=shared_experts_overlapping,
877
)
878
879
return self._maybe_combine(
vllm/model_executor/layers/fused_moe/runner/shared_experts.py
+39
-37
@@ -71,6 +71,11 @@ class SharedExperts(torch.nn.Module):
71
if self._stream is not None:
72
logger.debug_once("Enabled separate cuda stream for MoE shared_experts")
73
74
+ if self._stream is not None:
75
+ # One pair per DBO ubatch id.
76
+ self._input_ready_event = [torch.cuda.Event(), torch.cuda.Event()]
77
+ self._output_ready_event = [torch.cuda.Event(), torch.cuda.Event()]
78
+
79
# TODO(bnell): Hack for elastic_ep. Get rid of this
80
def _set_moe_config(self, new_moe_config: FusedMoEConfig):
81
self.moe_config = new_moe_config
@@ -96,6 +101,14 @@ class SharedExperts(torch.nn.Module):
101
and parallel_config.all2all_backend not in _EPLB_OVERLAP_SAFE_BACKENDS
102
) or parallel_config.use_fi_nvl_two_sided_kernels
103
104
+ @property
105
+ def _should_enable_stream_overlap_heuristic(self) -> bool:
106
+ # On ROCm, empirically it's shown that only DPA deployments benefit from
107
+ # multi-stream shared experts
108
+ if not current_platform.is_rocm():
109
+ return True
110
+ return self._moe_config.moe_parallel_config.dp_size > 1
111
+
112
def _determine_shared_experts_order(
113
self,
114
hidden_states: torch.Tensor,
@@ -107,10 +120,11 @@ class SharedExperts(torch.nn.Module):
120
return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED
121
122
should_run_shared_in_aux_stream = (
110
- current_platform.is_cuda()
123
+ current_platform.is_cuda_alike()
124
and self._stream is not None
125
and hidden_states.shape[0]
126
<= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
127
+ and self._should_enable_stream_overlap_heuristic
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