[ROCm][CI] Expand AITER W4A4 MoE Coverage (#52647)

Signed-off-by: Micah Williamson <micah.williamson@amd.com>

Author: micah-wil

tests/kernels/moe/test_ocp_mx_moe.py

@@ -125,7 +125,6 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
 
         # if model_case.model_id == "fxmarty/qwen_1.5-moe-a2.7b-mxfp4":
         #     llm.apply_model(check_model)
-
         output = llm.generate_greedy("Today I am in the French Alps and", max_tokens=20)
         assert output
 
@@ -187,6 +186,16 @@ def mxfp8_dequantize(x, scale):
     return x_float * scale
 
 
+def mxfp4_quant_dequant(x: torch.Tensor) -> torch.Tensor:
+    shape = x.shape
+    quantized, scale = dynamic_mxfp4_quant(x.to(torch.bfloat16).flatten(0, -2))
+    return (
+        upcast_from_mxfp(quantized.view(torch.uint8), scale, torch.bfloat16, axis=-1)
+        .reshape(shape)
+        .float()
+    )
+
+
 def reference_moe(
     roouting_logits,
     topk,
@@ -219,6 +228,8 @@ def reference_moe(
     expert_weights = torch.nn.functional.softmax(experts.values, dim=1)
     expert_indices = experts.indices
     t = hidden_states.clone()
+    if act_type == "mxfp4":
+        t = mxfp4_quant_dequant(t)
     # MLP #1
     mlp1_weight = w13[expert_indices, ...]
     mlp1_bias = bias13[expert_indices, ...]
@@ -244,6 +255,8 @@ def reference_moe(
             t.to(torch.bfloat16), is_sf_swizzled_layout=False
         )
         t = mxfp8_dequantize(t_quantized, t_scale)
+    elif act_type == "mxfp4":
+        t = mxfp4_quant_dequant(t)
     # MLP #2
     mlp2_weight = w2[expert_indices, ...]
     mlp2_bias = bias2[expert_indices, ...]
@@ -1392,6 +1405,14 @@ def test_trtllm_gen_mxfp8_block_scale_moe(
         "requires_aiter": True,
         "requires_gfx950": True,
     },
+    "AITER_MXFP4_MXFP4": {
+        "activation": "SILU",
+        "act_type": "mxfp4",
+        "rtol": 1.0,
+        "percent": 0.8,
+        "requires_aiter": True,
+        "requires_gfx950": True,
+    },
 }
 
 
@@ -1662,7 +1683,7 @@ class MockLayer:
         alpha=1.702 if activation == MoEActivation.SWIGLUOAI else 1.0,
         beta=1.0 if activation == MoEActivation.SWIGLUOAI else 0.0,
         limit=7.0 if activation == MoEActivation.SWIGLUOAI else None,
-        act_type="bf16",
+        act_type=str(config.get("act_type", "bf16")),
         activation=act_name,
         use_interleaved_layout=use_interleaved,
     )

tests/quantization/test_gfx950_moe.py

@@ -25,14 +25,32 @@
 
 ROCM_AVAILABLE = current_platform.is_rocm()
 ROCM_GFX950 = False
-ROCM_AITER_AVAILABLE = False
+ROCM_AITER_SUPPORTED = False
 
 if ROCM_AVAILABLE:
-    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
     from vllm.platforms.rocm import on_gfx950
 
     ROCM_GFX950 = on_gfx950()
-    ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_fused_moe_enabled()
+    ROCM_AITER_SUPPORTED = is_aiter_found_and_supported()
+
+
+def set_rocm_aiter(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None:
+    value = "1" if enabled else "0"
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER", value)
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", value)
+    monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", enabled)
+    monkeypatch.setattr(rocm_aiter_ops, "_FMOE_ENABLED", enabled)
+
+
+@pytest.fixture
+def enable_rocm_aiter(monkeypatch: pytest.MonkeyPatch):
+    set_rocm_aiter(monkeypatch, True)
+
+
+@pytest.fixture
+def disable_rocm_aiter(monkeypatch: pytest.MonkeyPatch):
+    set_rocm_aiter(monkeypatch, False)
 
 
 def _make_w4a4_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
@@ -68,8 +86,8 @@ def mxfp4_oracle_config():
 
 
 @pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
-@pytest.mark.skipif(not ROCM_AITER_AVAILABLE, reason="Requires AITER enabled")
-def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config):
+@pytest.mark.skipif(not ROCM_AITER_SUPPORTED, reason="Requires supported AITER")
+def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config, enable_rocm_aiter):
     """With AITER enabled + GFX950, W4A4 selects AITER_MXFP4_MXFP4."""
     config = _make_w4a4_moe_config()
     backend, experts_cls = select_mxfp4_moe_backend(
@@ -79,6 +97,19 @@ def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config):
     assert experts_cls is not None
 
 
+@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
+def test_w4a4_falls_back_without_aiter(
+    mxfp4_oracle_config,
+    disable_rocm_aiter,
+):
+    config = _make_w4a4_moe_config()
+    backend, experts_cls = select_mxfp4_moe_backend(
+        config, activation_key=kMxfp4Dynamic
+    )
+    assert backend == Mxfp4MoeBackend.EMULATION
+    assert experts_cls is not None
+
+
 @pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
 def test_w4a4_dispatches_to_emulation_with_moe_backend(mxfp4_oracle_config):
     """With --moe-backend emulation, W4A4 selects EMULATION."""