[ROCm][AITER][Bugfix] Disable emulation for MoE

tests/quantization/test_quark.py

@@ -385,3 +385,78 @@ def test_mxfp4_dequant_kernel_match_quark(
     out_torch = dq_mxfp4_torch(w_mxfp4, scale, float_dtype)
 
     assert torch.equal(out_hip, out_torch)
+
+
+def test_moe_emulation_w_mxfp4_a_mxfp4():
+    """w_mxfp4_a_mxfp4 on gfx950 (supports_mx + aiter) must not fall back
+    to emulation — it should use the native AITER CK path."""
+    from unittest.mock import MagicMock, patch
+
+    from vllm.model_executor.layers.fused_moe import FusedMoEConfig, MoEActivation
+    from vllm.model_executor.layers.fused_moe.config import (
+        FusedMoEParallelConfig,
+        RoutingMethodType,
+    )
+    from vllm.model_executor.layers.quantization.quark.quark_moe import (
+        QuarkOCP_MX_MoEMethod,
+    )
+
+    weight_config = {"dtype": "fp4", "qscheme": "per_group", "is_dynamic": False}
+    input_config = {"dtype": "fp4", "qscheme": "per_group", "is_dynamic": True}
+
+    parallel_config = FusedMoEParallelConfig(
+        tp_size=1,
+        pcp_size=1,
+        dp_size=1,
+        ep_size=1,
+        tp_rank=0,
+        pcp_rank=0,
+        dp_rank=0,
+        ep_rank=0,
+        sp_size=1,
+        use_ep=False,
+        all2all_backend="",
+        enable_eplb=False,
+    )
+    moe = FusedMoEConfig(
+        num_experts=8,
+        experts_per_token=2,
+        hidden_dim=256,
+        intermediate_size_per_partition=512,
+        num_local_experts=8,
+        num_logical_experts=8,
+        activation=MoEActivation.SILU,
+        device="gpu",
+        routing_method=RoutingMethodType.Default,
+        moe_parallel_config=parallel_config,
+        in_dtype=torch.bfloat16,
+    )
+
+    mock_vllm_config = MagicMock()
+
+    with (
+        patch(
+            "vllm.model_executor.layers.quantization.quark.quark_moe.current_platform"
+        ) as mock_platform,
+        patch(
+            "vllm.model_executor.layers.quantization.quark.quark_moe.rocm_aiter_ops"
+        ) as mock_aiter,
+        patch(
+            "vllm.model_executor.layers.quantization.quark.quark_moe"
+            ".get_current_vllm_config",
+            return_value=mock_vllm_config,
+        ),
+    ):
+        mock_platform.supports_mx.return_value = True
+        mock_aiter.is_fused_moe_enabled.return_value = True
+
+        method = QuarkOCP_MX_MoEMethod(
+            weight_config=weight_config,
+            input_config=input_config,
+            moe=moe,
+        )
+
+    assert method.ocp_mx_scheme == "w_mxfp4_a_mxfp4"
+    assert method.emulate is False, (
+        "w_mxfp4_a_mxfp4 on gfx950 (supports_mx + aiter) must not emulate"
+    )

vllm/model_executor/layers/quantization/quark/quark_moe.py

@@ -1026,11 +1026,11 @@ def __init__(
         )
 
         # TODO(aiter): extend once rocm_aiter_fused_experts gains dispatch
-        # for the other OCP MX schemes. Today its CK MoE kernel only has an
-        # entry for `w_mxfp4` (w4a16); mixed schemes like `w_mxfp4_a_mxfp6_*`
-        # fall through to QuantMethod.NO and raise "Unsupported kernel config
-        # for moe heuristic dispatch".
-        _AITER_NATIVE_OCP_MX_SCHEMES = ("w_mxfp4",)
+        # for the other OCP MX schemes. Today its CK MoE kernel has entries
+        # for `w_mxfp4` (w4a16) and `w_mxfp4_a_mxfp4`; mixed schemes like
+        # `w_mxfp4_a_mxfp6_*` fall through to QuantMethod.NO and raise
+        # "Unsupported kernel config for moe heuristic dispatch".
+        _AITER_NATIVE_OCP_MX_SCHEMES = ("w_mxfp4", "w_mxfp4_a_mxfp4")
         self.emulate = (
             not current_platform.supports_mx()
             or self.ocp_mx_scheme not in _AITER_NATIVE_OCP_MX_SCHEMES