diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 82cf6c365ee6..e9cb579a6e86 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -325,7 +325,7 @@ def _validate_rocm_aiter_fused_moe_shared_expert_args( return has_shared_expert -def _rocm_aiter_fused_moe_impl( +def _aiter_fused_moe( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, @@ -356,6 +356,7 @@ def _rocm_aiter_fused_moe_impl( shared_w2_scale: torch.Tensor | None = None, shared_expert_id: int = -1, q_dtype_a: torch.dtype | None = None, + **fused_moe_kwargs, ) -> torch.Tensor: has_shared_expert = _validate_rocm_aiter_fused_moe_shared_expert_args( shared_w1, @@ -395,6 +396,7 @@ def _rocm_aiter_fused_moe_impl( # is checked at config time, so this is only reached on an AITER # build new enough to accept it. extra_kwargs["quant_dtype_a"] = q_dtype_a + extra_kwargs.update(fused_moe_kwargs) return fused_moe( hidden_states, @@ -422,6 +424,41 @@ def _rocm_aiter_fused_moe_impl( ) +def _rocm_aiter_fused_moe_impl( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weight: torch.Tensor, + topk_ids: torch.Tensor, + expert_mask: torch.Tensor | None = None, + activation_method: int = 0, + quant_method: int = 0, + doweight_stage1: bool = False, + w1_scale: torch.Tensor | None = None, + w2_scale: torch.Tensor | None = None, + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, + num_local_tokens: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, + hidden_pad: int = 0, + intermediate_pad: int = 0, + gate_mode: str = "", + bias1: torch.Tensor | None = None, + bias2: torch.Tensor | None = None, + moe_sorting_dispatch_policy: int = 0, + swiglu_limit: float = 0.0, + beta: float | None = None, + linear_beta: float | None = None, + shared_w1: torch.Tensor | None = None, + shared_w2: torch.Tensor | None = None, + shared_w1_scale: torch.Tensor | None = None, + shared_w2_scale: torch.Tensor | None = None, + shared_expert_id: int = -1, + q_dtype_a: torch.dtype | None = None, +) -> torch.Tensor: + return _aiter_fused_moe(**locals()) + + def _rocm_aiter_fused_moe_fake( hidden_states: torch.Tensor, w1: torch.Tensor, @@ -2883,6 +2920,17 @@ def fused_moe_supports_quant_dtype_a(cls) -> bool: return "quant_dtype_a" in inspect.signature(fused_moe).parameters + @classmethod + @if_aiter_supported + @functools.cache + def fused_moe_supports_accumulate_output(cls) -> bool: + """Probe whether the installed aiter.fused_moe accepts `accumulate_output`.""" + import inspect + + from aiter.fused_moe import fused_moe + + return "output_ready_event" in inspect.signature(fused_moe).parameters + @staticmethod def _probe_dsv4_i384_fhmoe_capability(num_tokens: int) -> bool: """Probe AITER's CSV-backed DSV4 native-I384 FHMoE contract.""" @@ -3575,7 +3623,48 @@ def fused_moe( shared_w2_scale: torch.Tensor | None = None, shared_expert_id: int = -1, q_dtype_a: torch.dtype | None = None, + addend: torch.Tensor | None = None, + addend_ready: torch.cuda.Event | None = None, ) -> torch.Tensor: + """With `addend`, accumulates into it and returns it, waiting on + `addend_ready` only before stage 2. Eager only: the event can't cross + the custom-op schema.""" + if addend is not None: + return _aiter_fused_moe( + hidden_states, + w1, + w2, + topk_weight, + topk_ids, + expert_mask, + activation_method, + quant_method, + doweight_stage1, + w1_scale, + w2_scale, + a1_scale, + a2_scale, + num_local_tokens, + output_dtype, + hidden_pad, + intermediate_pad, + gate_mode, + bias1, + bias2, + moe_sorting_dispatch_policy, + swiglu_limit, + beta, + linear_beta, + shared_w1, + shared_w2, + shared_w1_scale, + shared_w2_scale, + shared_expert_id, + q_dtype_a, + output=addend, + accumulate_output=True, + output_ready_event=addend_ready, + ) return torch.ops.vllm.rocm_aiter_fused_moe( hidden_states, w1, diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 501dd9ba43ce..e4d041d7f9b8 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -291,6 +291,8 @@ def rocm_aiter_fused_experts( shared_w1_scale: torch.Tensor | None = None, shared_w2_scale: torch.Tensor | None = None, shared_expert_id: int = -1, + addend: torch.Tensor | None = None, + addend_ready: torch.cuda.Event | None = None, ) -> torch.Tensor: """ROCm AITER fused MoE expert computation.""" if quant_config is None: @@ -338,7 +340,7 @@ def rocm_aiter_fused_experts( "AITER tkw1 kernel does not support `num_local_tokens`" ) - return rocm_aiter_ops.asm_moe_tkw1( + result = rocm_aiter_ops.asm_moe_tkw1( hidden_states, w1, w2, @@ -353,6 +355,11 @@ def rocm_aiter_fused_experts( expert_mask=expert_mask, activation_method=activation_method, ) + if addend is None: + return result + if addend_ready is not None: + torch.cuda.current_stream().wait_event(addend_ready) + return addend.add_(result) else: quant_method = QuantMethod.NO.value @@ -475,6 +482,8 @@ def rocm_aiter_fused_experts( shared_w2_scale=shared_w2_scale, shared_expert_id=shared_expert_id, q_dtype_a=q_dtype_a, + addend=addend, + addend_ready=addend_ready, ) @@ -537,6 +546,9 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() + def supports_output_addend(self) -> bool: + return rocm_aiter_ops.fused_moe_supports_accumulate_output() + def workspace_shapes( self, M: int, @@ -576,6 +588,8 @@ def apply( workspace2: torch.Tensor, expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, + addend: torch.Tensor | None = None, + addend_ready: torch.cuda.Event | None = None, ): # TODO(rob): rocm_aiter_fused_experts uses self.quant_config's # a_scales for static quantization. Update this to fit better @@ -601,6 +615,8 @@ def apply( num_local_tokens=num_local_tokens, output_dtype=output.dtype, moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(), + addend=addend, + addend_ready=addend_ready, ) # avoid redundant copy when output is a view of the result if ( diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index accf57c59e14..0d4f5a7149d6 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -954,6 +954,10 @@ def activation( def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce: raise NotImplementedError + def supports_output_addend(self) -> bool: + """Whether `apply` can accumulate into `addend` after `addend_ready`.""" + return False + @abstractmethod def apply( self, @@ -1336,6 +1340,7 @@ def _fused_experts( apply_router_weight_on_input: bool, expert_tokens_meta: ExpertTokensMetadata | None, output_alias: torch.Tensor | None = None, + shared_experts: SharedExperts | None = None, ) -> torch.Tensor | UnfinalizedMoEOutput: _, M_full, N, K, top_k = self.fused_experts.moe_problem_size( a1q, w1, w2, topk_ids @@ -1388,6 +1393,16 @@ def _fused_experts( elif use_output_alias: fused_out = output_alias + addend_kwargs: dict[str, Any] = {} + if ( + shared_experts is not None + and not self.prepare_finalize.supports_async() + and self.fused_experts.supports_output_addend() + ): + claimed = shared_experts.claim_output(fused_out) + if claimed is not None: + addend_kwargs = dict(addend=claimed[0], addend_ready=claimed[1]) + unfinalized = self.fused_experts.apply( output=fused_out, hidden_states=a1q, @@ -1404,6 +1419,7 @@ def _fused_experts( workspace2=workspace2, expert_tokens_meta=expert_tokens_meta, apply_router_weight_on_input=apply_router_weight_on_input, + **addend_kwargs, ) # Experts that stopped after GEMM2 wrote nothing into `fused_out`; the @@ -1575,6 +1591,7 @@ def apply( apply_router_weight_on_input=apply_router_weight_on_input, expert_tokens_meta=expert_tokens_meta, output_alias=output, + shared_experts=shared_experts, ) if lora_ctx is not None: @@ -1727,6 +1744,13 @@ def __init__( self._post_init_setup() + # Resolved eagerly: the runner queries this while being traced. + self._supports_output_addend = ( + isinstance(self.impl, FusedMoEKernelModularImpl) + and not self.prepare_finalize.supports_async() + and self.impl.fused_experts.supports_output_addend() + ) + @property def can_overlap_shared_experts(self) -> bool: if isinstance(self.impl, FusedMoEKernelModularImpl): @@ -1772,6 +1796,9 @@ def output_is_reduced(self) -> bool: def supports_deferred_moe_finalize(self) -> bool: return self.prepare_finalize.supports_deferred_moe_finalize() + def supports_output_addend(self) -> bool: + return self._supports_output_addend + def apply_monolithic( self, hidden_states: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index 05ad62331c72..c2408697b641 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -272,9 +272,8 @@ def fused_topk_bias( num_expert_group=num_expert_group, topk_group=num_expert_group, need_renorm=renormalize, + routed_scaling_factor=routed_scaling_factor, ) - if routed_scaling_factor != 1.0: - topk_weights *= routed_scaling_factor return topk_weights, topk_ids if scoring_func == "sqrtsoftplus": 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 dd42c8332ffd..4db2a8d38e50 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -426,6 +426,25 @@ def _fused_output_is_reduced(self) -> bool: and self._quant_method.moe_kernel.output_is_reduced() ) + @property + def _fold_shared_output(self) -> bool: + """Whether the MoE op returns the shared + routed sum.""" + moe_kernel = self._quant_method.moe_kernel + return ( + self._shared_experts is not None + and moe_kernel is not None + and moe_kernel.supports_output_addend() + and self.routed_input_transform is None + and self.routed_output_transform is None + and self.routed_scaling_factor == 1.0 + and not self.do_naive_dispatch_combine + and self.moe_config.pcp_size == 1 + and not self._fused_output_is_reduced + and not self._quant_method.has_unpadded_output + and self.moe_config.hidden_dim_unpadded + in (None, self.moe_config.hidden_dim) + ) + def _maybe_reduce_shared_expert_output( self, shared_output: torch.Tensor | None, @@ -646,10 +665,16 @@ def _apply_quant_method( 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, - fused_out, - ) + if self._shared_experts is None: + return None, fused_out + + claimed = self._shared_experts.output_claimed + shared_output = self._shared_experts.output + if not self._shared_experts.fold_output: + return shared_output, fused_out + if not claimed: + fused_out = shared_output + cast(torch.Tensor, fused_out) + return None, fused_out def _sequence_parallel_context(self): """Return a context manager for sequence-parallel token @@ -735,7 +760,12 @@ def forward( ) ) - result = self._forward_entry( + forward_entry = ( + torch.ops.vllm.moe_forward + if self._fold_shared_output + else self._forward_entry + ) + result = forward_entry( hidden_states, router_logits, shared_experts_input, @@ -860,8 +890,7 @@ def _maybe_combine( ) hidden_states = get_pcp_group().reduce_scatter(hidden_states, dim=0) - if self.shared_experts is not None: - assert shared_output is not None + if shared_output is not None: return shared_output, hidden_states else: return hidden_states @@ -897,6 +926,7 @@ def _forward_impl( # before routed expert dispatch. shared_experts_overlapping = False if self._shared_experts is not None: + self._shared_experts.fold_output = self._fold_shared_output shared_experts_overlapping = self._shared_experts.maybe_forward_async( shared_experts_input ) 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 e9d80b75c46e..63f0b3a58e6d 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -53,6 +53,10 @@ def __init__( # index is always 0 and the second output list element is ignored. self.enable_dbo = enable_dbo self._output: list[torch.Tensor | None] = [None, None] + self._output_async = [False, False] + self._output_claimed = [False, False] + # Set by the runner; gates `claim_output`. + self.fold_output = False self._layer = layer self._moe_config = moe_config @@ -158,6 +162,7 @@ def maybe_forward_async(self, shared_experts_input: torch.Tensor) -> bool: self._input_ready_event[idx].wait(self._stream) self._output[idx] = self._layer(shared_experts_input) self._output_ready_event[idx].record(self._stream) + self._output_async[idx] = True return True def wait(self) -> None: @@ -169,11 +174,38 @@ def wait(self) -> None: def _output_idx(self) -> int: return dbo_current_ubatch_id() if self.enable_dbo else 0 + def claim_output( + self, like: torch.Tensor + ) -> tuple[torch.Tensor, torch.cuda.Event | None] | None: + """Return (output, ready event or None) for the routed experts to + accumulate into, or None if it can't be used in place of `like`.""" + idx = self._output_idx + output = self._output[idx] + if ( + not self.fold_output + or output is None + or output.shape != like.shape + or output.dtype != like.dtype + or output.device != like.device + or not output.is_contiguous() + ): + return None + self._output_claimed[idx] = True + event = self._output_ready_event[idx] if self._output_async[idx] else None + return output, event + + @property + def output_claimed(self) -> bool: + return self._output_claimed[self._output_idx] + @property def output(self) -> torch.Tensor: - assert self._output[self._output_idx] is not None - output = self._output[self._output_idx] - self._output[self._output_idx] = None + idx = self._output_idx + output = self._output[idx] + assert output is not None + self._output[idx] = None + self._output_async[idx] = False + self._output_claimed[idx] = False return output def forward( diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 189e55d2269c..afc6de2d6576 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1975,6 +1975,7 @@ def apply( global_num_experts=layer.global_num_experts, apply_router_weight_on_input=layer.apply_router_weight_on_input, expert_map=layer.expert_map, + shared_experts=shared_experts, shared_experts_input=shared_experts_input, ) diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 541e90d19197..fddc810f3342 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -475,8 +475,8 @@ def __init__( # trivial single group (num_expert_group=topk_group=1) reduces to plain # top-k while applying the sigmoid + bias correction and appending the # always-on shared expert; aiter applies the routed scaling internally. - # Every other path (vLLM top-k bias router, or no fusion) applies the - # routed scaling to the MoE output here. + # The vLLM router weights a fused shared expert by 1 / scale, so the + # scale goes on the MoE output; unfused, it is folded into top-k. self.experts = FusedMoEFactory( num_experts=config.num_local_experts, top_k=config.num_experts_per_tok, @@ -494,7 +494,9 @@ def __init__( swiglu_alpha=config.swiglu_alpha, swiglu_beta=config.swiglu_beta, routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_output=not self.use_aiter_moe_fse, + apply_routed_scale_to_output=( + self.is_fused_shared_expert_enabled and not self.use_aiter_moe_fse + ), router_logits_dtype=self.gate.out_dtype, shared_experts=self.shared_experts, n_shared_experts=(