[ROCm][Perf] Use flydsl moe with Minimax-M3 mxfp8 weights on gfx950 and implemented moe-backend selection (#46184)
Signed-off-by: Hongxia Yang <hongxia.yang@amd.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Tan Pin Siang <tanpinsiang@gmail.com>
Hongxia Yang committed
Jun 27, 2026 at 06:22 UTC
867fd5e8ed6b0bfcf84b24f82658c9fb698a6d35
5 files changed
+374
-4
tests/kernels/moe/test_mxfp8_aiter_backend_selection.py
new
+135
@@ -0,0 +1,135 @@
1
+# SPDX-License-Identifier: Apache-2.0
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
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.
9
+"""
10
+
11
+import dataclasses
12
+from unittest.mock import patch
13
+
14
+import pytest
15
+
16
+from vllm.platforms import current_platform
17
+
18
+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.experts.aiter_mxfp8_moe import ( # noqa: E402
23
+ AiterMxfp8Experts,
24
+)
25
+from vllm.model_executor.layers.fused_moe.modular_kernel import ( # noqa: E402
26
+ FusedMoEActivationFormat,
27
+)
28
+from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( # noqa: E402
29
+ Fp8MoeBackend,
30
+)
31
+from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( # noqa: E402
32
+ _BACKEND_NAME_MAP,
33
+ _SUPPORTED_BACKENDS,
34
+ _mxfp8_backend_to_kernel_cls,
35
+ _select_kernel_cls,
36
+)
37
+from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
38
+ kMxfp8Dynamic,
39
+ kMxfp8Static,
40
+)
41
+
42
+_AITER_MOD = "vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe"
43
+
44
+
45
+def _config(ep_size: int = 1):
46
+ cfg = make_dummy_moe_config(num_experts=128, experts_per_token=4, hidden_dim=6144)
47
+ if ep_size != 1:
48
+ cfg = dataclasses.replace(
49
+ cfg,
50
+ moe_parallel_config=dataclasses.replace(
51
+ cfg.moe_parallel_config, ep_size=ep_size, use_ep=True
52
+ ),
53
+ )
54
+ return cfg
55
+
56
+
57
+def _gfx950():
58
+ """Patch the platform so the device gate (gfx950 / MX) passes off-ROCm."""
59
+ return patch.multiple(
60
+ f"{_AITER_MOD}.current_platform",
61
+ is_rocm=lambda: True,
62
+ supports_mx=lambda: True,
63
+ )
64
+
65
+
66
+def _flydsl_installed(present: bool):
67
+ return patch(f"{_AITER_MOD}.is_aiter_mxfp8_moe_available", return_value=present)
68
+
69
+
70
+def test_aiter_mxfp8_registered():
71
+ """The FlyDSL backend is auto-selectable and reachable via --moe-backend aiter."""
72
+ assert Fp8MoeBackend.AITER_MXFP8 in _SUPPORTED_BACKENDS
73
+ assert _BACKEND_NAME_MAP["aiter"] is Fp8MoeBackend.AITER_MXFP8
74
+ assert _mxfp8_backend_to_kernel_cls(Fp8MoeBackend.AITER_MXFP8) == [
75
+ AiterMxfp8Experts
76
+ ]
77
+
78
+
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
+
85
+@pytest.mark.parametrize("ep_size", [1, 2])
86
+def test_ep_supported(ep_size):
87
+ """FlyDSL accepts both TP and EP: apply() forwards expert_map as expert_mask."""
88
+ assert (
89
+ AiterMxfp8Experts._supports_parallel_config(
90
+ _config(ep_size).moe_parallel_config
91
+ )
92
+ is True
93
+ )
94
+
95
+
96
+@pytest.mark.parametrize(
97
+ "present,ep_size,supported,reason_substr",
98
+ [
99
+ (True, 1, True, None), # gfx950 + flydsl + TP -> selectable
100
+ (True, 2, True, None), # gfx950 + flydsl + EP -> selectable (expert_mask)
101
+ (False, 1, False, "flydsl package"), # package missing -> clear reason
102
+ ],
103
+)
104
+def test_is_supported_config(present, ep_size, supported, reason_substr):
105
+ with _gfx950(), _flydsl_installed(present):
106
+ ok, reason = AiterMxfp8Experts.is_supported_config(
107
+ AiterMxfp8Experts,
108
+ _config(ep_size),
109
+ kMxfp8Static,
110
+ kMxfp8Dynamic,
111
+ FusedMoEActivationFormat.Standard,
112
+ )
113
+ assert ok is supported
114
+ if reason_substr is not None:
115
+ assert reason_substr in reason
116
+
117
+
118
+def test_explicit_moe_backend_aiter():
119
+ """--moe-backend aiter: returns FlyDSL when usable (TP or EP), else a clear
120
+ ValueError when the flydsl package is missing."""
121
+ with _gfx950(), _flydsl_installed(True):
122
+ assert (
123
+ _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1))
124
+ is AiterMxfp8Experts
125
+ )
126
+ assert (
127
+ _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(2))
128
+ is AiterMxfp8Experts
129
+ )
130
+ with (
131
+ _gfx950(),
132
+ _flydsl_installed(False),
133
+ pytest.raises(ValueError, match="flydsl package"),
134
+ ):
135
+ _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1))
vllm/_aiter_ops.py
+30
@@ -171,6 +171,7 @@ def _rocm_aiter_fused_moe_impl(
171
bias1: torch.Tensor | None = None,
172
bias2: torch.Tensor | None = None,
173
moe_sorting_dispatch_policy: int = 0,
174
+ swiglu_limit: float = 0.0,
175
) -> torch.Tensor:
176
from aiter import ActivationType, QuantType
177
from aiter.fused_moe import fused_moe
@@ -203,6 +204,7 @@ def _rocm_aiter_fused_moe_impl(
204
bias1=bias1,
205
bias2=bias2,
206
moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
207
+ swiglu_limit=swiglu_limit,
208
**extra_kwargs,
209
)
210
@@ -229,6 +231,7 @@ def _rocm_aiter_fused_moe_fake(
231
bias1: torch.Tensor | None = None,
232
bias2: torch.Tensor | None = None,
233
moe_sorting_dispatch_policy: int = 0,
234
+ swiglu_limit: float = 0.0,
235
) -> torch.Tensor:
236
if output_dtype is not None:
237
return torch.empty_like(hidden_states, dtype=output_dtype)
@@ -2201,6 +2204,7 @@ class rocm_aiter_ops:
2204
bias1: torch.Tensor | None = None,
2205
bias2: torch.Tensor | None = None,
2206
moe_sorting_dispatch_policy: int = 0,
2207
+ swiglu_limit: float = 0.0,
2208
) -> torch.Tensor:
2209
return torch.ops.vllm.rocm_aiter_fused_moe(
2210
hidden_states,
@@ -2224,6 +2228,7 @@ class rocm_aiter_ops:
2228
bias1,
2229
bias2,
2230
moe_sorting_dispatch_policy,
2231
+ swiglu_limit,
2232
)
2233
2234
@staticmethod
@@ -2678,6 +2683,31 @@ class rocm_aiter_ops:
2683
2684
return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors)
2685
2686
+ @staticmethod
2687
+ def shuffle_mxfp8_moe_weights(
2688
+ w13: torch.Tensor,
2689
+ w2: torch.Tensor,
2690
+ w13_scale: torch.Tensor,
2691
+ w2_scale: torch.Tensor,
2692
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
2693
+ """Preshuffle MXFP8 MoE weights + E8M0 scales into AITER's FlyDSL layout:
2694
+ gate/up-interleaved weights, interleaved scale for w13 (gate/up), plain
2695
+ scale for w2 (the interleaved variant is gate/up-only and misaligns w2).
2696
+ """
2697
+ from aiter.ops.shuffle import shuffle_scale, shuffle_weight
2698
+
2699
+ num_experts = w13.shape[0]
2700
+ w13 = shuffle_weight(w13, is_guinterleave=True, gate_up=True)
2701
+ w2 = shuffle_weight(w2, is_guinterleave=True, gate_up=False)
2702
+ w13_scale = shuffle_scale(
2703
+ w13_scale.reshape(-1, w13_scale.shape[-1]),
2704
+ num_experts,
2705
+ is_guinterleave=True,
2706
+ gate_up=True,
2707
+ )
2708
+ w2_scale = shuffle_scale(w2_scale.reshape(-1, w2_scale.shape[-1]))
2709
+ return w13, w2, w13_scale, w2_scale
2710
+
2711
@staticmethod
2712
def flash_attn_varlen_func(
2713
q: torch.Tensor,
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py
new
+166
@@ -0,0 +1,166 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""MXFP8 (1x32 block, E8M0) MoE via AITER's FlyDSL two-stage grouped GEMM
4
+(gfx950); alternative to ``Mxfp8NativeTritonExperts``. Routes through
5
+``aiter.fused_moe`` (per_1x32, gate_mode=INTERLEAVE); weights are preshuffled in
6
+``convert_to_fp8_moe_kernel_format``.
7
+"""
8
+
9
+import torch
10
+
11
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12
+from vllm.logger import init_logger
13
+from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
14
+ Mxfp8TritonExpertsBase,
15
+)
16
+from vllm.platforms import current_platform
17
+
18
+logger = init_logger(__name__)
19
+
20
+
21
+def is_aiter_mxfp8_moe_available() -> bool:
22
+ """True when the FlyDSL MXFP8 MoE can run here: gfx950, the ``flydsl``
23
+ package is importable, AND the installed aiter carries the mxfp8 FlyDSL
24
+ 2-stage support from ROCm/aiter#3811.
25
+
26
+ ``flydsl`` and ``aiter`` are separate packages, so ``is_flydsl_available()``
27
+ (flydsl pkg + arch) is necessary but not sufficient: an older aiter without
28
+ #3811 still ships the flydsl pkg and the ``aiter.ops.flydsl`` module but a
29
+ broken/missing ``per_1x32 + fp8`` 2-stage path. Without this extra gate a
30
+ nightly lacking #3811 would wrongly select FlyDSL instead of falling back to
31
+ the native Triton dot_scaled path. #3811 added no probe-able public symbol,
32
+ so detect the ``minimax_m3_mxfp8`` tuned config it shipped. Every check fails
33
+ closed (returns False -> triton dot_scaled), which is always safe."""
34
+ if not (current_platform.is_rocm() and current_platform.supports_mx()):
35
+ return False
36
+ try:
37
+ import os
38
+
39
+ import aiter
40
+ from aiter.ops.flydsl.utils import is_flydsl_available
41
+
42
+ if not is_flydsl_available():
43
+ return False
44
+ return os.path.exists(
45
+ os.path.join(
46
+ os.path.dirname(aiter.__file__),
47
+ "configs",
48
+ "model_configs",
49
+ "minimax_m3_mxfp8_tuned_fmoe.csv",
50
+ )
51
+ )
52
+ except Exception:
53
+ return False
54
+
55
+
56
+class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
57
+ """MXFP8 MoE through AITER's FlyDSL two-stage grouped GEMM (gfx950)."""
58
+
59
+ @property
60
+ def quant_dtype(self) -> torch.dtype | str | None:
61
+ return self.quant_config.quant_dtype
62
+
63
+ @property
64
+ def block_shape(self) -> list[int] | None:
65
+ return self.quant_config.block_shape
66
+
67
+ @property
68
+ def expects_unquantized_inputs(self) -> bool:
69
+ # aiter.fused_moe MXFP8-quantizes the activations internally.
70
+ return True
71
+
72
+ @staticmethod
73
+ def _supports_current_device() -> bool:
74
+ # Device capability only (gfx950 / MX-capable ROCm). The flydsl package
75
+ # check lives in is_supported_config so a missing package is reported
76
+ # distinctly from an unsupported device.
77
+ return current_platform.is_rocm() and current_platform.supports_mx()
78
+
79
+ @staticmethod
80
+ def _supports_parallel_config(moe_parallel_config) -> bool:
81
+ # Both TP (expert_map=None) and EP are supported: apply() forwards the
82
+ # expert_map as aiter's ``expert_mask`` (the per-rank local-expert
83
+ # selection), mirroring the native rocm_aiter_moe path.
84
+ return True
85
+
86
+ @staticmethod
87
+ def is_supported_config(
88
+ cls, moe_config, weight_key, activation_key, activation_format
89
+ ):
90
+ is_supported, reason = super().is_supported_config(
91
+ cls, moe_config, weight_key, activation_key, activation_format
92
+ )
93
+ # _supports_current_device() only gates on the device; surface a clear
94
+ # reason when the device is fine but the flydsl package is missing.
95
+ if is_supported and not is_aiter_mxfp8_moe_available():
96
+ return False, (
97
+ "kernel requires the aiter flydsl package, which is not installed"
98
+ )
99
+ return is_supported, reason
100
+
101
+ def apply(
102
+ self,
103
+ output: torch.Tensor,
104
+ hidden_states: torch.Tensor,
105
+ w1: torch.Tensor,
106
+ w2: torch.Tensor,
107
+ topk_weights: torch.Tensor,
108
+ topk_ids: torch.Tensor,
109
+ activation,
110
+ global_num_experts: int,
111
+ expert_map: torch.Tensor | None,
112
+ a1q_scale: torch.Tensor | None,
113
+ a2_scale: torch.Tensor | None,
114
+ workspace13: torch.Tensor,
115
+ workspace2: torch.Tensor,
116
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
117
+ apply_router_weight_on_input: bool,
118
+ ):
119
+ from aiter import ActivationType, QuantType
120
+ from aiter.ops.flydsl.moe_common import GateMode
121
+
122
+ from vllm._aiter_ops import rocm_aiter_ops
123
+
124
+ # Re-tag the preshuffled weights: replace_parameter drops the
125
+ # is_shuffled flag, without which aiter picks a broken CK kernel.
126
+ w1.is_shuffled = True
127
+ w2.is_shuffled = True
128
+
129
+ limit = self.quant_config.gemm1_clamp_limit
130
+ swiglu_limit = 0.0 if limit is None else float(limit)
131
+
132
+ # Under EP, aiter expects ``expert_mask`` as a 0/1 *local-expert* mask
133
+ # over global ids with a trailing fake-expert sentinel slot
134
+ # (shape ``[global_num_experts + 1]``), NOT vLLM's expert_map (a
135
+ # global->local index map with -1 for non-local). Convert it; aiter
136
+ # derives the global->local compaction from the mask itself. ``None``
137
+ # under pure TP.
138
+ if expert_map is not None:
139
+ local_mask = (expert_map >= 0).to(torch.int32)
140
+ expert_mask = torch.cat([local_mask, local_mask.new_zeros(1)])
141
+ else:
142
+ expert_mask = None
143
+
144
+ # Route through the graph-safe ``rocm_aiter_fused_moe`` custom op so the
145
+ # call is captured under HIP graphs / torch.compile (a direct
146
+ # ``aiter.fused_moe`` is opaque to the dispatcher). aiter requires FP32
147
+ # routing weights / INT32 ids.
148
+ out = rocm_aiter_ops.fused_moe(
149
+ hidden_states,
150
+ w1,
151
+ w2,
152
+ topk_weights.to(torch.float32),
153
+ topk_ids.to(torch.int32),
154
+ expert_mask=expert_mask,
155
+ activation_method=ActivationType.Swiglu.value,
156
+ quant_method=QuantType.per_1x32.value,
157
+ doweight_stage1=apply_router_weight_on_input,
158
+ w1_scale=self.w1_scale_val,
159
+ w2_scale=self.w2_scale_val,
160
+ a1_scale=None,
161
+ a2_scale=None,
162
+ gate_mode=GateMode.INTERLEAVE.value,
163
+ swiglu_limit=swiglu_limit,
164
+ output_dtype=output.dtype,
165
+ )
166
+ output.copy_(out.to(output.dtype))
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+8
-2
@@ -59,7 +59,9 @@ class Fp8MoeBackend(Enum):
59
# MXFP8 MoE via a Triton ``dot_scaled`` kernel that lowers to CDNA4
60
# (gfx950) native MX matrix-core ops. Weights stay in MXFP8 (no load-time
61
# format conversion); the FP8 values + E8M0 scales are consumed directly.
62
- NATIVE_MXFP8 = "NATIVE_MXFP8"
62
+ TRITON_MXFP8 = "TRITON_MXFP8"
63
+ # MXFP8 MoE via AITER (FlyDSL two-stage grouped GEMM) on gfx950.
64
+ AITER_MXFP8 = "AITER_MXFP8"
65
66
67
def _get_priority_backends(
@@ -423,6 +425,10 @@ def convert_to_fp8_moe_kernel_format(
425
)
426
elif fp8_backend == Fp8MoeBackend.AITER:
427
w13, w2 = rocm_aiter_ops.shuffle_weights(w13, w2)
428
+ elif fp8_backend == Fp8MoeBackend.AITER_MXFP8:
429
+ w13, w2, w13_scale, w2_scale = rocm_aiter_ops.shuffle_mxfp8_moe_weights(
430
+ w13, w2, w13_scale, w2_scale
431
+ )
432
elif fp8_backend == Fp8MoeBackend.MARLIN:
433
weight_block_size = getattr(layer, "weight_block_size", None)
434
if weight_block_size == [1, 32]:
@@ -484,7 +490,7 @@ def convert_to_fp8_moe_kernel_format(
490
# EMULATION dequantizes weights at runtime; NATIVE_MXFP8 consumes
491
# the MXFP8 weights as-is — neither needs a load-time layout change.
492
Fp8MoeBackend.EMULATION,
487
- Fp8MoeBackend.NATIVE_MXFP8,
493
+ Fp8MoeBackend.TRITON_MXFP8,
494
]:
495
raise ValueError(f"Unsupported FP8 MoE backend: {fp8_backend.value}")
496
vllm/model_executor/layers/fused_moe/oracle/mxfp8.py
+35
-2
@@ -21,6 +21,10 @@ _SUPPORTED_BACKENDS = (
21
Fp8MoeBackend.DEEPGEMM,
22
Fp8MoeBackend.MARLIN,
23
Fp8MoeBackend.XPU,
24
+ # AITER FlyDSL (gfx950): auto-picked by select_mxfp8_moe_backend when
25
+ # is_supported_config passes (gfx950 + flydsl installed + not EP). On other
26
+ # devices / no flydsl / EP it is skipped and native is used.
27
+ Fp8MoeBackend.AITER_MXFP8,
28
)
29
30
_BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = {
@@ -28,6 +32,8 @@ _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = {
32
"deep_gemm": Fp8MoeBackend.DEEPGEMM,
33
"marlin": Fp8MoeBackend.MARLIN,
34
"xpu": Fp8MoeBackend.XPU,
35
+ "aiter": Fp8MoeBackend.AITER_MXFP8,
36
+ "triton": Fp8MoeBackend.TRITON_MXFP8,
37
}
38
39
@@ -46,6 +52,27 @@ def _mxfp8_backend_to_kernel_cls(
52
)
53
54
return [DeepGemmExperts]
55
+ if backend == Fp8MoeBackend.AITER_MXFP8:
56
+ from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import (
57
+ AiterMxfp8Experts,
58
+ )
59
+
60
+ return [AiterMxfp8Experts]
61
+ if backend == Fp8MoeBackend.TRITON_MXFP8:
62
+ # Explicit ``--moe-backend triton``: the Triton mxfp8 path, i.e.
63
+ # dot_scaled on MX-capable HW (gfx950) and BF16 emulation otherwise.
64
+ # Mirrors the ROCm auto-fallback in ``_select_rocm_mxfp8_backend``.
65
+ if current_platform.supports_mx():
66
+ from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
67
+ Mxfp8NativeTritonExperts,
68
+ )
69
+
70
+ return [Mxfp8NativeTritonExperts]
71
+ from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
72
+ Mxfp8EmulationTritonExperts,
73
+ )
74
+
75
+ return [Mxfp8EmulationTritonExperts]
76
return backend_to_kernel_cls(backend)
77
78
@@ -77,7 +104,13 @@ def _select_kernel_cls(
104
105
106
def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
80
- """ROCm fallback when vendor MXFP8 backends are unavailable."""
107
+ """ROCm fallback when no auto-selected MXFP8 backend is available.
108
+
109
+ The aiter FlyDSL backend (``AITER_MXFP8``) is auto-picked earlier by
110
+ ``select_mxfp8_moe_backend`` via ``_SUPPORTED_BACKENDS`` when usable, or
111
+ explicitly via ``--moe-backend aiter``; this fallback handles the rest
112
+ (native dot_scaled on gfx950, else BF16 emulation).
113
+ """
114
115
if current_platform.supports_mx():
116
from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
@@ -85,7 +118,7 @@ def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts
118
)
119
120
logger.info_once("Using native CDNA4 (gfx950) MXFP8 dot_scaled MoE backend.")
88
- return Fp8MoeBackend.NATIVE_MXFP8, Mxfp8NativeTritonExperts
121
+ return Fp8MoeBackend.TRITON_MXFP8, Mxfp8NativeTritonExperts
122
123
from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
124
Mxfp8EmulationTritonExperts,