diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 1d75b7c7628a..d744da0b89b5 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -167,6 +167,7 @@ def _rocm_aiter_fused_moe_impl( 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, @@ -177,6 +178,10 @@ def _rocm_aiter_fused_moe_impl( activation = ActivationType(activation_method) quant_type = QuantType(quant_method) + extra_kwargs: dict = {} + if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode(): + extra_kwargs["gate_mode"] = gate_mode + return fused_moe( hidden_states, w1, @@ -198,6 +203,7 @@ def _rocm_aiter_fused_moe_impl( bias1=bias1, bias2=bias2, moe_sorting_dispatch_policy=moe_sorting_dispatch_policy, + **extra_kwargs, ) @@ -219,6 +225,7 @@ def _rocm_aiter_fused_moe_fake( 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, @@ -1804,6 +1811,21 @@ def are_gdn_triton_kernels_available(cls) -> bool: except (ImportError, ModuleNotFoundError): return False + @classmethod + @if_aiter_supported + @functools.cache + def fused_moe_supports_gate_mode(cls) -> bool: + """Probe whether the installed aiter.fused_moe accepts `gate_mode`. + + Added in https://github.com/ROCm/aiter/pull/3123 (>=0.1.14). + Builds with older AITER must omit this argument. + """ + import inspect + + from aiter.fused_moe import fused_moe + + return "gate_mode" in inspect.signature(fused_moe).parameters + @staticmethod @if_aiter_supported def register_ops_once() -> None: @@ -2172,6 +2194,7 @@ def fused_moe( 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, @@ -2194,6 +2217,7 @@ def fused_moe( output_dtype, hidden_pad, intermediate_pad, + gate_mode, bias1, bias2, moe_sorting_dispatch_policy, 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 5c2aa4556008..bd9b285fe74f 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 @@ -351,6 +351,21 @@ def rocm_aiter_fused_experts( intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1) ) + # https://github.com/ROCm/aiter/pull/3123 specialized the AITER stage1 GEMMs + # for interleaved vs separated gate and up weights. + # For gpt-oss i.e. use_mxfp4_w4a16=True, the weights are shuffled by + # `rocm_aiter_ops.shuffle_weight_a16w4` in `oracle/mxfp4.py`, + # which always sets `is_guinterleave=True`. + # Hence, we pass in GateMode.INTERLEAVE to match the weight shuffling. + gate_mode = "" + if quant_config.use_mxfp4_w4a16: + try: + from aiter.ops.flydsl.moe_common import GateMode + + gate_mode = GateMode.INTERLEAVE.value + except ImportError: + pass + return rocm_aiter_ops.fused_moe( hidden_states, w1, @@ -369,6 +384,7 @@ def rocm_aiter_fused_experts( output_dtype=output_dtype, hidden_pad=hidden_pad, intermediate_pad=intermediate_pad, + gate_mode=gate_mode, bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None, bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None, moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,