From 5369f7b7b89e48d396fa98a9d33f3fba654af8ee Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Wed, 29 Jul 2026 01:57:20 +0200 Subject: [PATCH] [MXFP8][ROCm] Fix MXFP8 MoE backend selection (#49747) Signed-off-by: Felix Marty --- .../moe/test_mxfp8_aiter_backend_selection.py | 44 ++++- tests/models/quantization/test_mxfp8.py | 166 ++++++++++++++++++ .../fused_moe/experts/aiter_mxfp8_moe.py | 27 +++ .../fused_moe/experts/mxfp8_emulation_moe.py | 8 +- .../fused_moe/experts/mxfp8_native_moe.py | 9 +- .../layers/fused_moe/oracle/mxfp8.py | 51 +----- 6 files changed, 244 insertions(+), 61 deletions(-) diff --git a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py index 7c2fdbabe29..aa1747f4153 100644 --- a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py +++ b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py @@ -19,9 +19,17 @@ if not current_platform.is_rocm(): pytest.skip("This test can only run on ROCm.", allow_module_level=True) from tests.kernels.moe.utils import make_dummy_moe_config # noqa: E402 +from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402 + MoEActivation, +) from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import ( # noqa: E402 + _AITER_SWIGLU_ALPHA, + _AITER_SWIGLU_BETA, AiterMxfp8Experts, ) +from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( # noqa: E402 + Mxfp8NativeTritonExperts, +) from vllm.model_executor.layers.fused_moe.modular_kernel import ( # noqa: E402 FusedMoEActivationFormat, ) @@ -33,6 +41,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( # noqa: E402 _SUPPORTED_BACKENDS, _mxfp8_backend_to_kernel_cls, _select_kernel_cls, + select_mxfp8_moe_backend, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 kMxfp8Dynamic, @@ -43,7 +52,17 @@ _AITER_MOD = "vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe" def _config(ep_size: int = 1): - cfg = make_dummy_moe_config(num_experts=128, experts_per_token=4, hidden_dim=6144) + # AiterMxfp8Experts hardcodes SwiGLU-OAI: match its required activation and + # alpha/beta so is_supported_config doesn't reject the config on those grounds. + cfg = make_dummy_moe_config( + num_experts=128, + experts_per_token=4, + hidden_dim=6144, + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ) + cfg = dataclasses.replace( + cfg, swiglu_alpha=_AITER_SWIGLU_ALPHA, swiglu_beta=_AITER_SWIGLU_BETA + ) if ep_size != 1: cfg = dataclasses.replace( cfg, @@ -76,12 +95,6 @@ def test_aiter_mxfp8_registered(): ] -def test_triton_selectable(): - assert _BACKEND_NAME_MAP["triton"] is Fp8MoeBackend.TRITON_MXFP8 - # Not auto-selected (only reachable explicitly), so FlyDSL still wins auto. - assert Fp8MoeBackend.TRITON_MXFP8 not in _SUPPORTED_BACKENDS - - @pytest.mark.parametrize("ep_size", [1, 2]) def test_ep_supported(ep_size): """FlyDSL accepts both TP and EP: apply() forwards expert_map as expert_mask.""" @@ -133,3 +146,20 @@ def test_explicit_moe_backend_aiter(): pytest.raises(ValueError, match="flydsl package"), ): _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1)) + + +def test_gfx950_picks_aiter(): + """Auto-select on real ROCm hardware with flydsl usable -> FlyDSL wins.""" + with _flydsl_installed(True): + backend, experts_cls = select_mxfp8_moe_backend(_config()) + assert backend is Fp8MoeBackend.AITER_MXFP8 + assert experts_cls is AiterMxfp8Experts + + +def test_gfx942_picks_triton(): + """flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton + dot_scaled backend wins instead.""" + with _flydsl_installed(False): + backend, experts_cls = select_mxfp8_moe_backend(_config()) + assert backend is Fp8MoeBackend.TRITON_MXFP8 + assert experts_cls is Mxfp8NativeTritonExperts diff --git a/tests/models/quantization/test_mxfp8.py b/tests/models/quantization/test_mxfp8.py index 7c250d11576..c12a72a09c0 100644 --- a/tests/models/quantization/test_mxfp8.py +++ b/tests/models/quantization/test_mxfp8.py @@ -17,8 +17,10 @@ diverse prompts from ``tests/prompts/example.txt``. """ import pytest +import torch from tests.quantization.utils import is_quant_method_supported +from vllm.platforms import current_platform from ..utils import check_logprobs_close @@ -81,6 +83,170 @@ def test_mxfp8_logprobs( ) +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_requires_swigluoai_activation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + ) + + with pytest.raises(ValueError, match="requires activation=swigluoai_uninterleave"): + select_mxfp8_moe_backend(config) + + +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_requires_swigluoai_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + ) + + with pytest.raises(ValueError, match="hardcodes SwiGLU-OAI"): + select_mxfp8_moe_backend(config) + + +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_accepts_swigluoai_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + swiglu_alpha=aiter_mxfp8_moe._AITER_SWIGLU_ALPHA, + swiglu_beta=aiter_mxfp8_moe._AITER_SWIGLU_BETA, + ) + + backend, experts_cls = select_mxfp8_moe_backend(config) + + assert backend == Fp8MoeBackend.AITER_MXFP8 + assert experts_cls is aiter_mxfp8_moe.AiterMxfp8Experts + + @pytest.mark.skipif( not is_quant_method_supported("mxfp8"), reason="mxfp8 is not supported on this GPU type (requires sm_100+).", diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index c5330f3b438..59115df13b5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -6,10 +6,13 @@ ``convert_to_fp8_moe_kernel_format``. """ +import math + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( Mxfp8TritonExpertsBase, ) @@ -17,6 +20,9 @@ from vllm.platforms import current_platform logger = init_logger(__name__) +_AITER_SWIGLU_ALPHA = 1.702 +_AITER_SWIGLU_BETA = 1.0 + def is_aiter_mxfp8_moe_available() -> bool: """True when the FlyDSL MXFP8 MoE can run here: gfx950, the ``flydsl`` @@ -93,6 +99,27 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase): return False, ( "kernel requires the aiter flydsl package, which is not installed" ) + if ( + is_supported + and moe_config.activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE + ): + return False, ( + "kernel hardcodes SwiGLU-OAI activation and requires " + f"activation={MoEActivation.SWIGLUOAI_UNINTERLEAVE.value}; " + f"got activation={moe_config.activation.value}" + ) + if is_supported and ( + moe_config.swiglu_alpha is None + or not math.isclose(float(moe_config.swiglu_alpha), _AITER_SWIGLU_ALPHA) + or moe_config.swiglu_beta is None + or not math.isclose(float(moe_config.swiglu_beta), _AITER_SWIGLU_BETA) + ): + return False, ( + "kernel hardcodes SwiGLU-OAI with " + f"alpha={_AITER_SWIGLU_ALPHA} and beta={_AITER_SWIGLU_BETA}; " + f"got swiglu_alpha={moe_config.swiglu_alpha} and " + f"swiglu_beta={moe_config.swiglu_beta}" + ) return is_supported, reason def apply( diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py index 71dd7634a69..ad6083251fb 100644 --- a/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py @@ -107,17 +107,13 @@ class Mxfp8EmulationTritonExperts(Mxfp8TritonExpertsBase): limit = self.quant_config.gemm1_clamp_limit if limit is None: raise ValueError("SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit") - alpha = self.quant_config.gemm1_alpha - alpha = 1.702 if alpha is None else float(alpha) - beta = self.quant_config.gemm1_beta - beta = 1.0 if beta is None else float(beta) apply_moe_activation( activation, output, input, clamp_limit=float(limit), - alpha=alpha, - beta=beta, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, ) return super().activation(activation, output, input) diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py index 9839756880a..e8c7dc96921 100644 --- a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py @@ -362,10 +362,7 @@ class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase): expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): - alpha = self.quant_config.gemm1_alpha - alpha = 1.702 if alpha is None else float(alpha) - beta = self.quant_config.gemm1_beta - beta = 1.0 if beta is None else float(beta) + # `self.gemm1_alpha` and `self.gemm1_beta`` are set by `TritonExperts.__init__`. limit = self.quant_config.gemm1_clamp_limit limit = None if limit is None else float(limit) out = fused_moe_mxfp8_native( @@ -376,8 +373,8 @@ class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase): self.w2_scale_val, topk_weights, topk_ids, - alpha=alpha, - beta=beta, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, limit=limit, global_num_experts=global_num_experts, expert_map=expert_map, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index b9086cfa48a..2a03eacea8e 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -12,10 +12,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kMxfp8Dynamic, kMxfp8Static, ) -from vllm.platforms import current_platform logger = init_logger(__name__) +# Ordered by priority. _SUPPORTED_BACKENDS = ( Fp8MoeBackend.FLASHINFER_TRTLLM, Fp8MoeBackend.DEEPGEMM, @@ -26,6 +26,8 @@ _SUPPORTED_BACKENDS = ( # devices / no flydsl / EP it is skipped and native is used. Fp8MoeBackend.AITER_MXFP8, Fp8MoeBackend.HUMMING, + Fp8MoeBackend.TRITON_MXFP8, + Fp8MoeBackend.EMULATION, ) _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { @@ -61,15 +63,12 @@ def _mxfp8_backend_to_kernel_cls( return [AiterMxfp8Experts] if backend == Fp8MoeBackend.TRITON_MXFP8: - # Explicit ``--moe-backend triton``: the Triton mxfp8 path, i.e. - # dot_scaled on MX-capable HW (gfx950) and BF16 emulation otherwise. - # Mirrors the ROCm auto-fallback in ``_select_rocm_mxfp8_backend``. - if current_platform.supports_mx(): - from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( - Mxfp8NativeTritonExperts, - ) + from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( + Mxfp8NativeTritonExperts, + ) - return [Mxfp8NativeTritonExperts] + return [Mxfp8NativeTritonExperts] + if backend == Fp8MoeBackend.EMULATION: from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( Mxfp8EmulationTritonExperts, ) @@ -105,35 +104,6 @@ def _select_kernel_cls( ) -def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: - """ROCm fallback when no auto-selected MXFP8 backend is available. - - The aiter FlyDSL backend (``AITER_MXFP8``) is auto-picked earlier by - ``select_mxfp8_moe_backend`` via ``_SUPPORTED_BACKENDS`` when usable, or - explicitly via ``--moe-backend aiter``; this fallback handles the rest - (native dot_scaled on gfx950, else BF16 emulation). - """ - - if current_platform.supports_mx(): - from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( - Mxfp8NativeTritonExperts, - ) - - logger.info_once("Using native CDNA4 (gfx950) MXFP8 dot_scaled MoE backend.") - return Fp8MoeBackend.TRITON_MXFP8, Mxfp8NativeTritonExperts - - from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( - Mxfp8EmulationTritonExperts, - ) - - logger.info_once( - "No native MXFP8 MoE backend available on this device; " - "MXFP8 weights will be dequantized to BF16 once at load time and the " - "MoE will run in BF16 (no per-step dequant)." - ) - return Fp8MoeBackend.EMULATION, Mxfp8EmulationTritonExperts - - def select_mxfp8_moe_backend( config: FusedMoEConfig, ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: @@ -167,8 +137,5 @@ def select_mxfp8_moe_backend( logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value) return backend, experts_cls - # simplify the logic for rocm, refactor later when more backends are supported - if current_platform.is_rocm(): - return _select_rocm_mxfp8_backend() - + # TODO: add debug log with reason. raise ValueError("No MXFP8 MoE backends available.")