[ROCm][Bugfix] Gate AITER MXFP8 MoE on the aiter enable flag (#57426)
Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Co-authored-by: Claude <noreply@anthropic.com>
Rohan Potdar committed
Sep 17, 2026 at 19:06 UTC
213c379fcafe5ac3bdc235a480e13140ee13c7b5
3 files changed
+24
-75
tests/kernels/moe/test_mxfp8_aiter_backend_selection.py
+22
-15
@@ -2,10 +2,11 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""MXFP8 MoE backend selection for the AITER FlyDSL kernel (gfx950).
4
5
-GPU-free: mocks the platform (gfx950) and the ``flydsl`` package check, then
5
+GPU-free: mocks the platform (gfx950) and the aiter fused-MoE enable flag, then
6
exercises the oracle so the FlyDSL backend is auto-picked when usable (including
7
under expert parallelism, since apply() forwards the expert_map as aiter's
8
-expert_mask) and skipped (native fallback) when the device/package is missing.
8
+expert_mask) and skipped (native fallback) when the device is unsupported or the
9
+aiter runtime is disabled.
10
"""
11
12
import dataclasses
@@ -48,6 +49,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa:
49
kMxfp8Static,
50
)
51
52
+from vllm._aiter_ops import rocm_aiter_ops # noqa: E402 # isort: skip
53
+
54
_AITER_MOD = "vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe"
55
56
@@ -82,8 +85,8 @@ def _gfx950():
85
)
86
87
85
-def _flydsl_installed(present: bool):
86
- return patch(f"{_AITER_MOD}.is_aiter_mxfp8_moe_available", return_value=present)
88
+def _aiter_moe_enabled(present: bool):
89
+ return patch.object(rocm_aiter_ops, "is_fused_moe_enabled", return_value=present)
90
91
92
def test_aiter_mxfp8_registered():
@@ -109,13 +112,18 @@ def test_ep_supported(ep_size):
112
@pytest.mark.parametrize(
113
"present,ep_size,supported,reason_substr",
114
[
112
- (True, 1, True, None), # gfx950 + flydsl + TP -> selectable
113
- (True, 2, True, None), # gfx950 + flydsl + EP -> selectable (expert_mask)
114
- (False, 1, False, "flydsl package"), # package missing -> clear reason
115
+ (True, 1, True, None), # gfx950 + aiter MoE + TP -> selectable
116
+ (True, 2, True, None), # gfx950 + aiter MoE + EP -> selectable (expert_mask)
117
+ (
118
+ False,
119
+ 1,
120
+ False,
121
+ "does not support current device",
122
+ ), # disabled -> not selected
123
],
124
)
125
def test_is_supported_config(present, ep_size, supported, reason_substr):
118
- with _gfx950(), _flydsl_installed(present):
126
+ with _gfx950(), _aiter_moe_enabled(present):
127
ok, reason = AiterMxfp8Experts.is_supported_config(
128
AiterMxfp8Experts,
129
_config(ep_size),
@@ -130,8 +138,8 @@ def test_is_supported_config(present, ep_size, supported, reason_substr):
138
139
def test_explicit_moe_backend_aiter():
140
"""--moe-backend aiter: returns FlyDSL when usable (TP or EP), else a clear
133
- ValueError when the flydsl package is missing."""
134
- with _gfx950(), _flydsl_installed(True):
141
+ ValueError when the aiter MoE runtime is disabled."""
142
+ with _gfx950(), _aiter_moe_enabled(True):
143
assert (
144
_select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1))
145
is AiterMxfp8Experts
@@ -142,18 +150,17 @@ def test_explicit_moe_backend_aiter():
150
)
151
with (
152
_gfx950(),
145
- _flydsl_installed(False),
146
- pytest.raises(ValueError, match="flydsl package"),
153
+ _aiter_moe_enabled(False),
154
+ pytest.raises(ValueError, match="does not support current device"),
155
):
156
_select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1))
157
158
159
def test_gfx950_picks_aiter():
152
- """Auto-select on real ROCm hardware with flydsl usable -> FlyDSL wins."""
153
- # NOTE: Fp8MoeBackend.AITER_MXFP8 does not require VLLM_ROCM_USE_AITER=1
160
+ """Auto-select on real ROCm hardware with aiter MoE enabled -> FlyDSL wins."""
161
with (
162
patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=True),
156
- _flydsl_installed(True),
163
+ _aiter_moe_enabled(True),
164
):
165
backend, experts_cls = select_mxfp8_moe_backend(_config())
166
assert backend is Fp8MoeBackend.AITER_MXFP8
tests/models/quantization/test_mxfp8.py
-15
@@ -111,11 +111,6 @@ def test_mxfp8_aiter_requires_swigluoai_activation(
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
- )
114
115
config = FusedMoEConfig(
116
num_experts=8,
@@ -164,11 +159,6 @@ def test_mxfp8_aiter_requires_swigluoai_params(
159
"_supports_current_device",
160
staticmethod(lambda: True),
161
)
167
- monkeypatch.setattr(
168
- aiter_mxfp8_moe,
169
- "is_aiter_mxfp8_moe_available",
170
- lambda: True,
171
- )
162
163
config = FusedMoEConfig(
164
num_experts=8,
@@ -218,11 +208,6 @@ def test_mxfp8_aiter_accepts_swigluoai_params(
208
"_supports_current_device",
209
staticmethod(lambda: True),
210
)
221
- monkeypatch.setattr(
222
- aiter_mxfp8_moe,
223
- "is_aiter_mxfp8_moe_available",
224
- lambda: True,
225
- )
211
212
config = FusedMoEConfig(
213
num_experts=8,
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py
+2
-45
@@ -11,6 +11,7 @@ import math
11
import torch
12
13
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14
+from vllm._aiter_ops import rocm_aiter_ops
15
from vllm.logger import init_logger
16
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
17
from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
@@ -24,41 +25,6 @@ _AITER_SWIGLU_ALPHA = 1.702
25
_AITER_SWIGLU_BETA = 1.0
26
27
27
-def is_aiter_mxfp8_moe_available() -> bool:
28
- """True when the FlyDSL MXFP8 MoE can run here: gfx950, the ``flydsl``
29
- package is importable, AND the installed aiter carries the mxfp8 FlyDSL
30
- 2-stage support from ROCm/aiter#3811.
31
-
32
- ``flydsl`` and ``aiter`` are separate packages, so ``is_flydsl_available()``
33
- (flydsl pkg + arch) is necessary but not sufficient: an older aiter without
34
- #3811 still ships the flydsl pkg and the ``aiter.ops.flydsl`` module but a
35
- broken/missing ``per_1x32 + fp8`` 2-stage path. Without this extra gate a
36
- nightly lacking #3811 would wrongly select FlyDSL instead of falling back to
37
- the native Triton dot_scaled path. #3811 added no probe-able public symbol,
38
- so detect the ``minimax_m3_mxfp8`` tuned config it shipped. Every check fails
39
- closed (returns False -> triton dot_scaled), which is always safe."""
40
- if not (current_platform.is_rocm() and current_platform.supports_mx()):
41
- return False
42
- try:
43
- import os
44
-
45
- import aiter
46
- from aiter.ops.flydsl.utils import is_flydsl_available
47
-
48
- if not is_flydsl_available():
49
- return False
50
- return os.path.exists(
51
- os.path.join(
52
- os.path.dirname(aiter.__file__),
53
- "configs",
54
- "model_configs",
55
- "minimax_m3_mxfp8_tuned_fmoe.csv",
56
- )
57
- )
58
- except Exception:
59
- return False
60
-
61
-
28
class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
29
"""MXFP8 MoE through AITER's FlyDSL two-stage grouped GEMM (gfx950)."""
30
@@ -79,10 +45,7 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
45
46
@staticmethod
47
def _supports_current_device() -> bool:
82
- # Device capability only (gfx950 / MX-capable ROCm). The flydsl package
83
- # check lives in is_supported_config so a missing package is reported
84
- # distinctly from an unsupported device.
85
- return current_platform.is_rocm() and current_platform.supports_mx()
48
+ return current_platform.supports_mx() and rocm_aiter_ops.is_fused_moe_enabled()
49
50
@staticmethod
51
def _supports_parallel_config(moe_parallel_config) -> bool:
@@ -95,12 +58,6 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
58
is_supported, reason = super().is_supported_config(
59
cls, moe_config, weight_key, activation_key, activation_format
60
)
98
- # _supports_current_device() only gates on the device; surface a clear
99
- # reason when the device is fine but the flydsl package is missing.
100
- if is_supported and not is_aiter_mxfp8_moe_available():
101
- return False, (
102
- "kernel requires the aiter flydsl package, which is not installed"
103
- )
61
if (
62
is_supported
63
and moe_config.activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE