diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 66ad630e68..2f072d14b4 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -283,11 +283,24 @@ class MoERunner(MoERunnerInterface): self._shared_experts: SharedExperts | None = None if shared_experts is not None: can_overlap = lambda: self._quant_method.mk_can_overlap_shared_experts + # When unquantized, shared expert inputs alias the hidden states, + # which can lead to race condition in multi-stream mode. Quantized + # routed experts copy the input into a fresh buffer first, breaking + # the alias, so overlap is safe. Only observed on ROCm. + routed_input_is_quantized = lambda: ( + self.routed_experts.quant_method.moe_quant_config is not None + and self.routed_experts.quant_method.moe_quant_config.quant_dtype + is not None + ) + is_multistream_safe = lambda: ( + not current_platform.is_rocm() or routed_input_is_quantized() + ) self._shared_experts = SharedExperts( shared_experts, moe_config=moe_config, enable_dbo=enable_dbo, mk_can_overlap_shared_experts=can_overlap, + is_multistream_safe=is_multistream_safe, ) # Needed for string -> MoERunner layer lookup in custom ops. @@ -582,12 +595,17 @@ class MoERunner(MoERunnerInterface): 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 | UnfinalizedMoEOutput]: """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 @@ -617,10 +635,9 @@ class MoERunner(MoERunnerInterface): shared_experts_input=shared_experts_input, ) - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.MULTI_STREAM_OVERLAPPED, - ) + if shared_experts_overlapping: + assert self._shared_experts is not None + self._shared_experts.wait() return ( self._shared_experts.output if self._shared_experts is not None else None, @@ -642,18 +659,6 @@ class MoERunner(MoERunnerInterface): 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, @@ -878,8 +883,13 @@ class MoERunner(MoERunnerInterface): # TODO(bnell): this can be removed after MK migration is complete. self.routed_experts._ensure_moe_quant_config_init() - # Sync aux and main stream for shared expert multi-stream overlap. - self._maybe_sync_shared_experts_stream(shared_experts_input) + # 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 + ) # If the Runner holds the gate, apply it after the stream sync, # so it can run overlapped with the @@ -905,6 +915,7 @@ class MoERunner(MoERunnerInterface): router_logits=router_logits, shared_experts_input=shared_experts_input, input_ids=input_ids, + shared_experts_overlapping=shared_experts_overlapping, ) return self._maybe_combine( diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index 17faeffd29..5b9587c853 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -43,6 +43,7 @@ class SharedExperts(torch.nn.Module): moe_config: FusedMoEConfig, enable_dbo: bool, mk_can_overlap_shared_experts: Callable[[], bool], + is_multistream_safe: Callable[[], bool], ): super().__init__() @@ -57,6 +58,10 @@ class SharedExperts(torch.nn.Module): self._mk_can_overlap_shared_experts = mk_can_overlap_shared_experts + # Might not be safe to run multi-stream mode if routed and shared experts + # alias the same inputs + self._is_multistream_safe = is_multistream_safe + # Allow disabling of the separate shared experts stream for # debug purposes. # TODO: Remove this after more extensive testings with TP/DP @@ -65,11 +70,12 @@ class SharedExperts(torch.nn.Module): logger.debug_once("Disabling MoE shared_experts cuda stream") self._stream = None else: - # TODO(rob): enable shared expert overlap with non-cuda-alike. - # aux_stream() returns None on non-cuda-alike platforms. self._stream = aux_stream() if self._stream is not None: logger.debug_once("Enabled separate cuda stream for MoE shared_experts") + # One pair per DBO ubatch id to sync aux and main stream. + 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): @@ -106,11 +112,19 @@ class SharedExperts(torch.nn.Module): if self._mk_can_overlap_shared_experts(): return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED + # On ROCm, empirically only DP-only deployments benefit from the overlap. + overlap_is_beneficial = not current_platform.is_rocm() or ( + self._moe_config.moe_parallel_config.dp_size > 1 + and self._moe_config.moe_parallel_config.tp_size == 1 + ) + should_run_shared_in_aux_stream = ( - current_platform.is_cuda() + current_platform.is_cuda_alike() and self._stream is not None and hidden_states.shape[0] <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD + and overlap_is_beneficial + and self._is_multistream_safe() ) if should_run_shared_in_aux_stream: @@ -118,38 +132,31 @@ class SharedExperts(torch.nn.Module): else: return SharedExpertsOrder.NO_OVERLAP - def maybe_sync_shared_experts_stream( - self, - shared_experts_input: torch.Tensor, - ): - experts_order = self._determine_shared_experts_order(shared_experts_input) + def maybe_forward_async(self, shared_experts_input: torch.Tensor) -> bool: + """Enqueue shared experts on the aux stream without waiting for them. - if experts_order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED: - assert self._stream is not None - - # 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. + 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()) with torch.cuda.stream(self._stream): - output = self._layer(shared_experts_input) - current_stream().wait_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 - return output + 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()) @property def _output_idx(self) -> int: @@ -174,11 +181,6 @@ class SharedExperts(torch.nn.Module): assert self._output[self._output_idx] is None - 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) + self._output[self._output_idx] = self._layer(shared_experts_input) assert self._output[self._output_idx] is not None