Revert "[Perf][ROCm] Dual-stream decode with hipgraphs" (#52024)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Simon Danielsson committed
Aug 13, 2026 at 06:38 UTC
3d204dfdaaf09d67d49c7855630ef949754e0f8f
2 files changed
+55
-55
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+18
-16
@@ -573,17 +573,12 @@ 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,
576
) -> tuple[torch.Tensor | None, torch.Tensor]:
577
"""Run expert routing and the fused MoE kernel via the quant method.
578
579
Orchestrates shared expert execution (before/after), expert selection
580
via the router, and the actual fused MoE computation. Returns
581
(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.
582
"""
583
self._maybe_apply_shared_experts(
584
shared_experts_input, SharedExpertsOrder.NO_OVERLAP
@@ -613,9 +608,10 @@ class MoERunner(MoERunnerInterface):
608
shared_experts_input=shared_experts_input,
609
)
610
616
- if shared_experts_overlapping:
617
- assert self._shared_experts is not None
618
- self._shared_experts.wait()
611
+ self._maybe_apply_shared_experts(
612
+ shared_experts_input,
613
+ SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
614
+ )
615
616
return (
617
self._shared_experts.output if self._shared_experts is not None else None,
@@ -637,6 +633,18 @@ class MoERunner(MoERunnerInterface):
633
else nullcontext()
634
)
635
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
+
648
def _maybe_add_zero_expert_output(
649
self,
650
result: torch.Tensor,
@@ -841,13 +849,8 @@ class MoERunner(MoERunnerInterface):
849
# TODO(bnell): this can be removed after MK migration is complete.
850
self.routed_experts._ensure_moe_quant_config_init()
851
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
- )
852
+ # Sync aux and main stream for shared expert multi-stream overlap.
853
+ self._maybe_sync_shared_experts_stream(shared_experts_input)
854
855
# If the Runner holds the gate, apply it after the stream sync,
856
# so it can run overlapped with the
@@ -873,7 +876,6 @@ class MoERunner(MoERunnerInterface):
876
router_logits=router_logits,
877
shared_experts_input=shared_experts_input,
878
input_ids=input_ids,
876
- shared_experts_overlapping=shared_experts_overlapping,
879
)
880
881
return self._maybe_combine(
vllm/model_executor/layers/fused_moe/runner/shared_experts.py
+37
-39
@@ -71,11 +71,6 @@ 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
-
74
# TODO(bnell): Hack for elastic_ep. Get rid of this
75
def _set_moe_config(self, new_moe_config: FusedMoEConfig):
76
self.moe_config = new_moe_config
@@ -101,14 +96,6 @@ class SharedExperts(torch.nn.Module):
96
and parallel_config.all2all_backend not in _EPLB_OVERLAP_SAFE_BACKENDS
97
) or parallel_config.use_fi_nvl_two_sided_kernels
98
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
-
99
def _determine_shared_experts_order(
100
self,
101
hidden_states: torch.Tensor,
@@ -120,11 +107,10 @@ class SharedExperts(torch.nn.Module):
107
return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED
108
109
should_run_shared_in_aux_stream = (
123
- current_platform.is_cuda_alike()
110
+ current_platform.is_cuda()
111
and self._stream is not None
112
and hidden_states.shape[0]
113
<= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
127
- and self._should_enable_stream_overlap_heuristic
114
)
115
116
if should_run_shared_in_aux_stream:
@@ -132,31 +118,38 @@ class SharedExperts(torch.nn.Module):
118
else:
119
return SharedExpertsOrder.NO_OVERLAP
120
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())
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.
148
with torch.cuda.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
149
+ output = self._layer(shared_experts_input)
150
+ current_stream().wait_stream(self._stream)
151
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())
152
+ return output
153
154
@property
155
def _output_idx(self) -> int:
@@ -181,6 +174,11 @@ class SharedExperts(torch.nn.Module):
174
175
assert self._output[self._output_idx] is None
176
184
- self._output[self._output_idx] = self._layer(shared_experts_input)
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)
183
184
assert self._output[self._output_idx] is not None