Upstream sync 17/N: merge 3d204dfdaa Revert "[Perf][ROCm] Dual-stream decode with hipgraphs" (#52024) (conflict)

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -573,17 +573,12 @@ def _apply_quant_method(
         router_logits: torch.Tensor,
         shared_experts_input: torch.Tensor | None,
         input_ids: torch.Tensor | None = None,
-        shared_experts_overlapping: bool = False,
     ) -> tuple[torch.Tensor | None, torch.Tensor]:
         """Run expert routing and the fused MoE kernel via the quant method.
 
         Orchestrates shared expert execution (before/after), expert selection
         via the router, and the actual fused MoE computation. Returns
         (shared_expert_output, fused_expert_output).
-
-        `shared_experts_overlapping` should be True only if using multi-stream
-        overlap. Then the shared expert was already launched in a separate
-        stream, so the results only have to be awaited here.
         """
         self._maybe_apply_shared_experts(
             shared_experts_input, SharedExpertsOrder.NO_OVERLAP
@@ -613,9 +608,10 @@ def _apply_quant_method(
                 shared_experts_input=shared_experts_input,
             )
 
-        if shared_experts_overlapping:
-            assert self._shared_experts is not None
-            self._shared_experts.wait()
+        self._maybe_apply_shared_experts(
+            shared_experts_input,
+            SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
+        )
 
         return (
             self._shared_experts.output if self._shared_experts is not None else None,
@@ -637,6 +633,18 @@ def _sequence_parallel_context(self):
             else nullcontext()
         )
 
+    def _maybe_sync_shared_experts_stream(
+        self,
+        shared_experts_input: torch.Tensor | None,
+    ):
+        # If router/gate provided, then apply it here.
+        # (Note: This code runs only when "overlapped mode" is on to allow
+        #        parallel execution of shared experts with the RoutedExperts via
+        #        separate cuda stream)
+        if self._shared_experts is not None:
+            assert shared_experts_input is not None
+            self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
+
     def _maybe_add_zero_expert_output(
         self,
         result: torch.Tensor,
@@ -841,13 +849,8 @@ def _forward_impl(
         # TODO(bnell): this can be removed after MK migration is complete.
         self.routed_experts._ensure_moe_quant_config_init()
 
-        # If using multi-stream overlap for shared experts, we must launch it
-        # before routed expert dispatch.
-        shared_experts_overlapping = False
-        if self._shared_experts is not None:
-            shared_experts_overlapping = self._shared_experts.maybe_forward_async(
-                shared_experts_input
-            )
+        # Sync aux and main stream for shared expert multi-stream overlap.
+        self._maybe_sync_shared_experts_stream(shared_experts_input)
 
         # If the Runner holds the gate, apply it after the stream sync,
         # so it can run overlapped with the
@@ -873,7 +876,6 @@ def _forward_impl(
                 router_logits=router_logits,
                 shared_experts_input=shared_experts_input,
                 input_ids=input_ids,
-                shared_experts_overlapping=shared_experts_overlapping,
             )
 
             return self._maybe_combine(

vllm/model_executor/layers/fused_moe/runner/shared_experts.py

@@ -88,11 +88,6 @@ def __init__(
             if on_gfx11():
                 self._stream_token_threshold = sys.maxsize
 
-        if self._stream is not None:
-            # One pair per DBO ubatch id.
-            self._input_ready_event = [torch.cuda.Event(), torch.cuda.Event()]
-            self._output_ready_event = [torch.cuda.Event(), torch.cuda.Event()]
-
     # TODO(bnell): Hack for elastic_ep. Get rid of this
     def _set_moe_config(self, new_moe_config: FusedMoEConfig):
         self.moe_config = new_moe_config
@@ -118,23 +113,6 @@ def _disable_shared_experts_overlap(self) -> bool:
             and parallel_config.all2all_backend not in _EPLB_OVERLAP_SAFE_BACKENDS
         ) or parallel_config.use_fi_nvl_two_sided_kernels
 
-    @property
-    def _should_enable_stream_overlap_heuristic(self) -> bool:
-        # On ROCm, empirically it's shown that only DPA deployments benefit from
-        # multi-stream shared experts
-        if not current_platform.is_rocm():
-            return True
-
-        # gfx11 (RDNA3) is the exception: the overlap is net-positive at every
-        # measured prefill size there, single-GPU included, so it is not gated
-        # on data parallelism. See the token-threshold default in __init__.
-        from vllm.platforms.rocm import on_gfx11
-
-        if on_gfx11():
-            return True
-
-        return self._moe_config.moe_parallel_config.dp_size > 1
-
     def _determine_shared_experts_order(
         self,
         hidden_states: torch.Tensor,
@@ -149,39 +127,45 @@ def _determine_shared_experts_order(
             current_platform.is_cuda_alike()
             and self._stream is not None
             and hidden_states.shape[0] <= self._stream_token_threshold
-            and self._should_enable_stream_overlap_heuristic
         )
 
         if should_run_shared_in_aux_stream:
             return SharedExpertsOrder.MULTI_STREAM_OVERLAPPED
         else:
             return SharedExpertsOrder.NO_OVERLAP
 
-    def maybe_forward_async(self, shared_experts_input: torch.Tensor) -> bool:
-        """Enqueue shared experts on the aux stream without waiting for them.
+    def maybe_sync_shared_experts_stream(
+        self,
+        shared_experts_input: torch.Tensor,
+    ):
+        experts_order = self._determine_shared_experts_order(shared_experts_input)
+
+        if experts_order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED:
+            assert self._stream is not None
 
-        Returns true if the shared experts were enqueued, false otherwise. Call
-        `wait` to wait for the shared experts to finish if this returns true.
-        """
-        if (
-            self._determine_shared_experts_order(shared_experts_input)
-            != SharedExpertsOrder.MULTI_STREAM_OVERLAPPED
-        ):
-            return False
-        assert self._stream is not None
-        idx = self._output_idx
-        assert self._output[idx] is None
-        self._input_ready_event[idx].record(current_stream())
+            # Record that the clone will be used by shared_experts_stream
+            # to avoid gc issue from deallocation of hidden_states_clone
+            # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501
+            # NOTE: We don't need shared_output.record_stream(current_stream())
+            # because we synch the streams before using shared_output.
+            shared_experts_input.record_stream(self._stream)
+
+            # Mark sync start point for the aux stream since we will
+            # run in parallel with router/gate.
+            self._stream.wait_stream(current_stream())
+
+    def _run_in_aux_stream(
+        self,
+        shared_experts_input: torch.Tensor,
+    ) -> torch.Tensor:
+        # TODO: assert that maybe_sync_shared_experts_stream has been called.
+
+        # Run shared experts in parallel on a separate stream.
         with torch.cuda.stream(self._stream):
-            self._input_ready_event[idx].wait(self._stream)
-            self._output[idx] = self._layer(shared_experts_input)
-            self._output_ready_event[idx].record(self._stream)
-        return True
+            output = self._layer(shared_experts_input)
+        current_stream().wait_stream(self._stream)
 
-    def wait(self) -> None:
-        """Block the main stream until `maybe_forward_async` output is ready."""
-        assert self._stream is not None
-        self._output_ready_event[self._output_idx].wait(current_stream())
+        return output
 
     @property
     def _output_idx(self) -> int:
@@ -206,6 +190,11 @@ def forward(
 
         assert self._output[self._output_idx] is None
 
-        self._output[self._output_idx] = self._layer(shared_experts_input)
+        if order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED:
+            self._output[self._output_idx] = self._run_in_aux_stream(
+                shared_experts_input
+            )
+        else:
+            self._output[self._output_idx] = self._layer(shared_experts_input)
 
         assert self._output[self._output_idx] is not None