[ROCm][CI] Expand AITER W4A4 MoE Coverage (#52647)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Micah Williamson committed
Aug 17, 2026 at 15:46 UTC
8878ebd8fdfe286eb7cf00ccbb6919bc7a40b1dd
2 files changed
+59
-7
tests/kernels/moe/test_ocp_mx_moe.py
+23
-2
@@ -125,7 +125,6 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
125
126
# if model_case.model_id == "fxmarty/qwen_1.5-moe-a2.7b-mxfp4":
127
# llm.apply_model(check_model)
128
-
128
output = llm.generate_greedy("Today I am in the French Alps and", max_tokens=20)
129
assert output
130
@@ -187,6 +186,16 @@ def mxfp8_dequantize(x, scale):
186
return x_float * scale
187
188
189
+def mxfp4_quant_dequant(x: torch.Tensor) -> torch.Tensor:
190
+ shape = x.shape
191
+ quantized, scale = dynamic_mxfp4_quant(x.to(torch.bfloat16).flatten(0, -2))
192
+ return (
193
+ upcast_from_mxfp(quantized.view(torch.uint8), scale, torch.bfloat16, axis=-1)
194
+ .reshape(shape)
195
+ .float()
196
+ )
197
+
198
+
199
def reference_moe(
200
roouting_logits,
201
topk,
@@ -219,6 +228,8 @@ def reference_moe(
228
expert_weights = torch.nn.functional.softmax(experts.values, dim=1)
229
expert_indices = experts.indices
230
t = hidden_states.clone()
231
+ if act_type == "mxfp4":
232
+ t = mxfp4_quant_dequant(t)
233
# MLP #1
234
mlp1_weight = w13[expert_indices, ...]
235
mlp1_bias = bias13[expert_indices, ...]
@@ -244,6 +255,8 @@ def reference_moe(
255
t.to(torch.bfloat16), is_sf_swizzled_layout=False
256
)
257
t = mxfp8_dequantize(t_quantized, t_scale)
258
+ elif act_type == "mxfp4":
259
+ t = mxfp4_quant_dequant(t)
260
# MLP #2
261
mlp2_weight = w2[expert_indices, ...]
262
mlp2_bias = bias2[expert_indices, ...]
@@ -1392,6 +1405,14 @@ ROCM_BACKEND_CONFIGS = {
1405
"requires_aiter": True,
1406
"requires_gfx950": True,
1407
},
1408
+ "AITER_MXFP4_MXFP4": {
1409
+ "activation": "SILU",
1410
+ "act_type": "mxfp4",
1411
+ "rtol": 1.0,
1412
+ "percent": 0.8,
1413
+ "requires_aiter": True,
1414
+ "requires_gfx950": True,
1415
+ },
1416
}
1417
1418
@@ -1662,7 +1683,7 @@ def test_rocm_mxfp4_moe_oracle(
1683
alpha=1.702 if activation == MoEActivation.SWIGLUOAI else 1.0,
1684
beta=1.0 if activation == MoEActivation.SWIGLUOAI else 0.0,
1685
limit=7.0 if activation == MoEActivation.SWIGLUOAI else None,
1665
- act_type="bf16",
1686
+ act_type=str(config.get("act_type", "bf16")),
1687
activation=act_name,
1688
use_interleaved_layout=use_interleaved,
1689
)
tests/quantization/test_gfx950_moe.py
+36
-5
@@ -25,14 +25,32 @@ from vllm.platforms import current_platform
25
26
ROCM_AVAILABLE = current_platform.is_rocm()
27
ROCM_GFX950 = False
28
-ROCM_AITER_AVAILABLE = False
28
+ROCM_AITER_SUPPORTED = False
29
30
if ROCM_AVAILABLE:
31
- from vllm._aiter_ops import rocm_aiter_ops
31
+ from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
32
from vllm.platforms.rocm import on_gfx950
33
34
ROCM_GFX950 = on_gfx950()
35
- ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_fused_moe_enabled()
35
+ ROCM_AITER_SUPPORTED = is_aiter_found_and_supported()
36
+
37
+
38
+def set_rocm_aiter(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None:
39
+ value = "1" if enabled else "0"
40
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", value)
41
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", value)
42
+ monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", enabled)
43
+ monkeypatch.setattr(rocm_aiter_ops, "_FMOE_ENABLED", enabled)
44
+
45
+
46
+@pytest.fixture
47
+def enable_rocm_aiter(monkeypatch: pytest.MonkeyPatch):
48
+ set_rocm_aiter(monkeypatch, True)
49
+
50
+
51
+@pytest.fixture
52
+def disable_rocm_aiter(monkeypatch: pytest.MonkeyPatch):
53
+ set_rocm_aiter(monkeypatch, False)
54
55
56
def _make_w4a4_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
@@ -68,8 +86,8 @@ def mxfp4_oracle_config():
86
87
88
@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
71
-@pytest.mark.skipif(not ROCM_AITER_AVAILABLE, reason="Requires AITER enabled")
72
-def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config):
89
+@pytest.mark.skipif(not ROCM_AITER_SUPPORTED, reason="Requires supported AITER")
90
+def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config, enable_rocm_aiter):
91
"""With AITER enabled + GFX950, W4A4 selects AITER_MXFP4_MXFP4."""
92
config = _make_w4a4_moe_config()
93
backend, experts_cls = select_mxfp4_moe_backend(
@@ -79,6 +97,19 @@ def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config):
97
assert experts_cls is not None
98
99
100
+@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
101
+def test_w4a4_falls_back_without_aiter(
102
+ mxfp4_oracle_config,
103
+ disable_rocm_aiter,
104
+):
105
+ config = _make_w4a4_moe_config()
106
+ backend, experts_cls = select_mxfp4_moe_backend(
107
+ config, activation_key=kMxfp4Dynamic
108
+ )
109
+ assert backend == Mxfp4MoeBackend.EMULATION
110
+ assert experts_cls is not None
111
+
112
+
113
@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
114
def test_w4a4_dispatches_to_emulation_with_moe_backend(mxfp4_oracle_config):
115
"""With --moe-backend emulation, W4A4 selects EMULATION."""