@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5369f7b7b8

[MXFP8][ROCm] Fix MXFP8 MoE backend selection (#49747)

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

fxmarty-amd committed Jul 29, 2026 at 01:57 UTC 5369f7b7b89e48d396fa98a9d33f3fba654af8ee
6 files changed +244 -61
tests/kernels/moe/test_mxfp8_aiter_backend_selection.py
+37 -7
@@ -19,9 +19,17 @@ if not current_platform.is_rocm():
19 pytest.skip("This test can only run on ROCm.", allow_module_level=True)
20
21 from tests.kernels.moe.utils import make_dummy_moe_config # noqa: E402
22 +from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402
23 + MoEActivation,
24 +)
25 from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import ( # noqa: E402
26 + _AITER_SWIGLU_ALPHA,
27 + _AITER_SWIGLU_BETA,
28 AiterMxfp8Experts,
29 )
30 +from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( # noqa: E402
31 + Mxfp8NativeTritonExperts,
32 +)
33 from vllm.model_executor.layers.fused_moe.modular_kernel import ( # noqa: E402
34 FusedMoEActivationFormat,
35 )
@@ -33,6 +41,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( # noqa: E402
41 _SUPPORTED_BACKENDS,
42 _mxfp8_backend_to_kernel_cls,
43 _select_kernel_cls,
44 + select_mxfp8_moe_backend,
45 )
46 from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
47 kMxfp8Dynamic,
@@ -43,7 +52,17 @@ _AITER_MOD = "vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe"
52
53
54 def _config(ep_size: int = 1):
46 - cfg = make_dummy_moe_config(num_experts=128, experts_per_token=4, hidden_dim=6144)
55 + # AiterMxfp8Experts hardcodes SwiGLU-OAI: match its required activation and
56 + # alpha/beta so is_supported_config doesn't reject the config on those grounds.
57 + cfg = make_dummy_moe_config(
58 + num_experts=128,
59 + experts_per_token=4,
60 + hidden_dim=6144,
61 + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
62 + )
63 + cfg = dataclasses.replace(
64 + cfg, swiglu_alpha=_AITER_SWIGLU_ALPHA, swiglu_beta=_AITER_SWIGLU_BETA
65 + )
66 if ep_size != 1:
67 cfg = dataclasses.replace(
68 cfg,
@@ -76,12 +95,6 @@ def test_aiter_mxfp8_registered():
95 ]
96
97
79 -def test_triton_selectable():
80 - assert _BACKEND_NAME_MAP["triton"] is Fp8MoeBackend.TRITON_MXFP8
81 - # Not auto-selected (only reachable explicitly), so FlyDSL still wins auto.
82 - assert Fp8MoeBackend.TRITON_MXFP8 not in _SUPPORTED_BACKENDS
83 -
84 -
98 @pytest.mark.parametrize("ep_size", [1, 2])
99 def test_ep_supported(ep_size):
100 """FlyDSL accepts both TP and EP: apply() forwards expert_map as expert_mask."""
@@ -133,3 +146,20 @@ def test_explicit_moe_backend_aiter():
146 pytest.raises(ValueError, match="flydsl package"),
147 ):
148 _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1))
149 +
150 +
151 +def test_gfx950_picks_aiter():
152 + """Auto-select on real ROCm hardware with flydsl usable -> FlyDSL wins."""
153 + with _flydsl_installed(True):
154 + backend, experts_cls = select_mxfp8_moe_backend(_config())
155 + assert backend is Fp8MoeBackend.AITER_MXFP8
156 + assert experts_cls is AiterMxfp8Experts
157 +
158 +
159 +def test_gfx942_picks_triton():
160 + """flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton
161 + dot_scaled backend wins instead."""
162 + with _flydsl_installed(False):
163 + backend, experts_cls = select_mxfp8_moe_backend(_config())
164 + assert backend is Fp8MoeBackend.TRITON_MXFP8
165 + assert experts_cls is Mxfp8NativeTritonExperts
tests/models/quantization/test_mxfp8.py
+166
@@ -17,8 +17,10 @@ diverse prompts from ``tests/prompts/example.txt``.
17 """
18
19 import pytest
20 +import torch
21
22 from tests.quantization.utils import is_quant_method_supported
23 +from vllm.platforms import current_platform
24
25 from ..utils import check_logprobs_close
26
@@ -81,6 +83,170 @@ def test_mxfp8_logprobs(
83 )
84
85
86 +@pytest.mark.skipif(
87 + not is_quant_method_supported("mxfp8"),
88 + reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
89 +)
90 +@pytest.mark.skipif(
91 + not current_platform.is_rocm(),
92 + reason="AITER MXFP8 MoE backend is ROCm-only.",
93 +)
94 +@pytest.mark.quant_model
95 +def test_mxfp8_aiter_requires_swigluoai_activation(
96 + monkeypatch: pytest.MonkeyPatch,
97 +) -> None:
98 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
99 + from vllm.model_executor.layers.fused_moe.config import (
100 + FusedMoEConfig,
101 + FusedMoEParallelConfig,
102 + RoutingMethodType,
103 + )
104 + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe
105 + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
106 + select_mxfp8_moe_backend,
107 + )
108 +
109 + monkeypatch.setattr(
110 + aiter_mxfp8_moe.AiterMxfp8Experts,
111 + "_supports_current_device",
112 + staticmethod(lambda: True),
113 + )
114 + monkeypatch.setattr(
115 + aiter_mxfp8_moe,
116 + "is_aiter_mxfp8_moe_available",
117 + lambda: True,
118 + )
119 +
120 + config = FusedMoEConfig(
121 + num_experts=8,
122 + experts_per_token=2,
123 + hidden_dim=256,
124 + intermediate_size=256,
125 + num_local_experts=8,
126 + num_logical_experts=8,
127 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
128 + activation=MoEActivation.SILU,
129 + in_dtype=torch.bfloat16,
130 + device="cuda",
131 + routing_method=RoutingMethodType.Renormalize,
132 + moe_backend="aiter",
133 + )
134 +
135 + with pytest.raises(ValueError, match="requires activation=swigluoai_uninterleave"):
136 + select_mxfp8_moe_backend(config)
137 +
138 +
139 +@pytest.mark.skipif(
140 + not is_quant_method_supported("mxfp8"),
141 + reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
142 +)
143 +@pytest.mark.skipif(
144 + not current_platform.is_rocm(),
145 + reason="AITER MXFP8 MoE backend is ROCm-only.",
146 +)
147 +@pytest.mark.quant_model
148 +def test_mxfp8_aiter_requires_swigluoai_params(
149 + monkeypatch: pytest.MonkeyPatch,
150 +) -> None:
151 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
152 + from vllm.model_executor.layers.fused_moe.config import (
153 + FusedMoEConfig,
154 + FusedMoEParallelConfig,
155 + RoutingMethodType,
156 + )
157 + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe
158 + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
159 + select_mxfp8_moe_backend,
160 + )
161 +
162 + monkeypatch.setattr(
163 + aiter_mxfp8_moe.AiterMxfp8Experts,
164 + "_supports_current_device",
165 + staticmethod(lambda: True),
166 + )
167 + monkeypatch.setattr(
168 + aiter_mxfp8_moe,
169 + "is_aiter_mxfp8_moe_available",
170 + lambda: True,
171 + )
172 +
173 + config = FusedMoEConfig(
174 + num_experts=8,
175 + experts_per_token=2,
176 + hidden_dim=256,
177 + intermediate_size=256,
178 + num_local_experts=8,
179 + num_logical_experts=8,
180 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
181 + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
182 + in_dtype=torch.bfloat16,
183 + device="cuda",
184 + routing_method=RoutingMethodType.Renormalize,
185 + moe_backend="aiter",
186 + )
187 +
188 + with pytest.raises(ValueError, match="hardcodes SwiGLU-OAI"):
189 + select_mxfp8_moe_backend(config)
190 +
191 +
192 +@pytest.mark.skipif(
193 + not is_quant_method_supported("mxfp8"),
194 + reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
195 +)
196 +@pytest.mark.skipif(
197 + not current_platform.is_rocm(),
198 + reason="AITER MXFP8 MoE backend is ROCm-only.",
199 +)
200 +@pytest.mark.quant_model
201 +def test_mxfp8_aiter_accepts_swigluoai_params(
202 + monkeypatch: pytest.MonkeyPatch,
203 +) -> None:
204 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
205 + from vllm.model_executor.layers.fused_moe.config import (
206 + FusedMoEConfig,
207 + FusedMoEParallelConfig,
208 + RoutingMethodType,
209 + )
210 + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe
211 + from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
212 + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
213 + select_mxfp8_moe_backend,
214 + )
215 +
216 + monkeypatch.setattr(
217 + aiter_mxfp8_moe.AiterMxfp8Experts,
218 + "_supports_current_device",
219 + staticmethod(lambda: True),
220 + )
221 + monkeypatch.setattr(
222 + aiter_mxfp8_moe,
223 + "is_aiter_mxfp8_moe_available",
224 + lambda: True,
225 + )
226 +
227 + config = FusedMoEConfig(
228 + num_experts=8,
229 + experts_per_token=2,
230 + hidden_dim=256,
231 + intermediate_size=256,
232 + num_local_experts=8,
233 + num_logical_experts=8,
234 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
235 + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
236 + in_dtype=torch.bfloat16,
237 + device="cuda",
238 + routing_method=RoutingMethodType.Renormalize,
239 + moe_backend="aiter",
240 + swiglu_alpha=aiter_mxfp8_moe._AITER_SWIGLU_ALPHA,
241 + swiglu_beta=aiter_mxfp8_moe._AITER_SWIGLU_BETA,
242 + )
243 +
244 + backend, experts_cls = select_mxfp8_moe_backend(config)
245 +
246 + assert backend == Fp8MoeBackend.AITER_MXFP8
247 + assert experts_cls is aiter_mxfp8_moe.AiterMxfp8Experts
248 +
249 +
250 @pytest.mark.skipif(
251 not is_quant_method_supported("mxfp8"),
252 reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py
+27
@@ -6,10 +6,13 @@
6 ``convert_to_fp8_moe_kernel_format``.
7 """
8
9 +import math
10 +
11 import torch
12
13 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14 from vllm.logger import init_logger
15 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
16 from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
17 Mxfp8TritonExpertsBase,
18 )
@@ -17,6 +20,9 @@ from vllm.platforms import current_platform
20
21 logger = init_logger(__name__)
22
23 +_AITER_SWIGLU_ALPHA = 1.702
24 +_AITER_SWIGLU_BETA = 1.0
25 +
26
27 def is_aiter_mxfp8_moe_available() -> bool:
28 """True when the FlyDSL MXFP8 MoE can run here: gfx950, the ``flydsl``
@@ -93,6 +99,27 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
99 return False, (
100 "kernel requires the aiter flydsl package, which is not installed"
101 )
102 + if (
103 + is_supported
104 + and moe_config.activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE
105 + ):
106 + return False, (
107 + "kernel hardcodes SwiGLU-OAI activation and requires "
108 + f"activation={MoEActivation.SWIGLUOAI_UNINTERLEAVE.value}; "
109 + f"got activation={moe_config.activation.value}"
110 + )
111 + if is_supported and (
112 + moe_config.swiglu_alpha is None
113 + or not math.isclose(float(moe_config.swiglu_alpha), _AITER_SWIGLU_ALPHA)
114 + or moe_config.swiglu_beta is None
115 + or not math.isclose(float(moe_config.swiglu_beta), _AITER_SWIGLU_BETA)
116 + ):
117 + return False, (
118 + "kernel hardcodes SwiGLU-OAI with "
119 + f"alpha={_AITER_SWIGLU_ALPHA} and beta={_AITER_SWIGLU_BETA}; "
120 + f"got swiglu_alpha={moe_config.swiglu_alpha} and "
121 + f"swiglu_beta={moe_config.swiglu_beta}"
122 + )
123 return is_supported, reason
124
125 def apply(
vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py
+2 -6
@@ -107,17 +107,13 @@ class Mxfp8EmulationTritonExperts(Mxfp8TritonExpertsBase):
107 limit = self.quant_config.gemm1_clamp_limit
108 if limit is None:
109 raise ValueError("SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit")
110 - alpha = self.quant_config.gemm1_alpha
111 - alpha = 1.702 if alpha is None else float(alpha)
112 - beta = self.quant_config.gemm1_beta
113 - beta = 1.0 if beta is None else float(beta)
110 apply_moe_activation(
111 activation,
112 output,
113 input,
114 clamp_limit=float(limit),
119 - alpha=alpha,
120 - beta=beta,
115 + alpha=self.gemm1_alpha,
116 + beta=self.gemm1_beta,
117 )
118 return
119 super().activation(activation, output, input)
vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py
+3 -6
@@ -362,10 +362,7 @@ class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase):
362 expert_tokens_meta: mk.ExpertTokensMetadata | None,
363 apply_router_weight_on_input: bool,
364 ):
365 - alpha = self.quant_config.gemm1_alpha
366 - alpha = 1.702 if alpha is None else float(alpha)
367 - beta = self.quant_config.gemm1_beta
368 - beta = 1.0 if beta is None else float(beta)
365 + # `self.gemm1_alpha` and `self.gemm1_beta`` are set by `TritonExperts.__init__`.
366 limit = self.quant_config.gemm1_clamp_limit
367 limit = None if limit is None else float(limit)
368 out = fused_moe_mxfp8_native(
@@ -376,8 +373,8 @@ class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase):
373 self.w2_scale_val,
374 topk_weights,
375 topk_ids,
379 - alpha=alpha,
380 - beta=beta,
376 + alpha=self.gemm1_alpha,
377 + beta=self.gemm1_beta,
378 limit=limit,
379 global_num_experts=global_num_experts,
380 expert_map=expert_map,
vllm/model_executor/layers/fused_moe/oracle/mxfp8.py
+9 -42
@@ -12,10 +12,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
12 kMxfp8Dynamic,
13 kMxfp8Static,
14 )
15 -from vllm.platforms import current_platform
15
16 logger = init_logger(__name__)
17
18 +# Ordered by priority.
19 _SUPPORTED_BACKENDS = (
20 Fp8MoeBackend.FLASHINFER_TRTLLM,
21 Fp8MoeBackend.DEEPGEMM,
@@ -26,6 +26,8 @@ _SUPPORTED_BACKENDS = (
26 # devices / no flydsl / EP it is skipped and native is used.
27 Fp8MoeBackend.AITER_MXFP8,
28 Fp8MoeBackend.HUMMING,
29 + Fp8MoeBackend.TRITON_MXFP8,
30 + Fp8MoeBackend.EMULATION,
31 )
32
33 _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = {
@@ -61,15 +63,12 @@ def _mxfp8_backend_to_kernel_cls(
63
64 return [AiterMxfp8Experts]
65 if backend == Fp8MoeBackend.TRITON_MXFP8:
64 - # Explicit ``--moe-backend triton``: the Triton mxfp8 path, i.e.
65 - # dot_scaled on MX-capable HW (gfx950) and BF16 emulation otherwise.
66 - # Mirrors the ROCm auto-fallback in ``_select_rocm_mxfp8_backend``.
67 - if current_platform.supports_mx():
68 - from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
69 - Mxfp8NativeTritonExperts,
70 - )
66 + from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
67 + Mxfp8NativeTritonExperts,
68 + )
69
72 - return [Mxfp8NativeTritonExperts]
70 + return [Mxfp8NativeTritonExperts]
71 + if backend == Fp8MoeBackend.EMULATION:
72 from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
73 Mxfp8EmulationTritonExperts,
74 )
@@ -105,35 +104,6 @@ def _select_kernel_cls(
104 )
105
106
108 -def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
109 - """ROCm fallback when no auto-selected MXFP8 backend is available.
110 -
111 - The aiter FlyDSL backend (``AITER_MXFP8``) is auto-picked earlier by
112 - ``select_mxfp8_moe_backend`` via ``_SUPPORTED_BACKENDS`` when usable, or
113 - explicitly via ``--moe-backend aiter``; this fallback handles the rest
114 - (native dot_scaled on gfx950, else BF16 emulation).
115 - """
116 -
117 - if current_platform.supports_mx():
118 - from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
119 - Mxfp8NativeTritonExperts,
120 - )
121 -
122 - logger.info_once("Using native CDNA4 (gfx950) MXFP8 dot_scaled MoE backend.")
123 - return Fp8MoeBackend.TRITON_MXFP8, Mxfp8NativeTritonExperts
124 -
125 - from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
126 - Mxfp8EmulationTritonExperts,
127 - )
128 -
129 - logger.info_once(
130 - "No native MXFP8 MoE backend available on this device; "
131 - "MXFP8 weights will be dequantized to BF16 once at load time and the "
132 - "MoE will run in BF16 (no per-step dequant)."
133 - )
134 - return Fp8MoeBackend.EMULATION, Mxfp8EmulationTritonExperts
135 -
136 -
107 def select_mxfp8_moe_backend(
108 config: FusedMoEConfig,
109 ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
@@ -167,8 +137,5 @@ def select_mxfp8_moe_backend(
137 logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value)
138 return backend, experts_cls
139
170 - # simplify the logic for rocm, refactor later when more backends are supported
171 - if current_platform.is_rocm():
172 - return _select_rocm_mxfp8_backend()
173 -
140 + # TODO: add debug log with reason.
141 raise ValueError("No MXFP8 MoE backends available.")