[Bugfix][ROCm] Reject unsupported EP for monolithic AITER MXFP4 MoE (#57866)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Andreas Karatzas committed
Sep 21, 2026 at 15:30 UTC
12972ae40bece6f9b9898fadd0c50d6a317df6fd
3 files changed
+57
-2
tests/kernels/moe/test_ocp_mx_moe.py
+55
@@ -1868,6 +1868,61 @@ def test_select_mxfp4_moe_backend_raises_with_unsupported_reasons(
1868
mxfp4_oracle.select_mxfp4_moe_backend(moe_config)
1869
1870
1871
+@pytest.mark.parametrize("backend_name", ["AITER_TRITON_MXFP4_BF16", "AITER_MXFP4_FP8"])
1872
+@pytest.mark.parametrize("use_ep", [False, True])
1873
+def test_aiter_mxfp4_monolithic_rejects_expert_parallel(
1874
+ monkeypatch: pytest.MonkeyPatch, backend_name: str, use_ep: bool
1875
+):
1876
+ """Global routing cannot index local expert weights, even without all-to-all."""
1877
+ from tests.kernels.moe.utils import make_dummy_moe_config
1878
+ from vllm.model_executor.layers.fused_moe.activation import MoEActivation
1879
+ from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
1880
+ from vllm.model_executor.layers.fused_moe.modular_kernel import (
1881
+ FusedMoEActivationFormat,
1882
+ )
1883
+ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
1884
+ Mxfp4MoeBackend,
1885
+ backend_to_kernel_cls,
1886
+ )
1887
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
1888
+ kFp8StaticTensorSym,
1889
+ kMxfp4Static,
1890
+ )
1891
+
1892
+ experts_cls = backend_to_kernel_cls(Mxfp4MoeBackend[backend_name])[0]
1893
+ monkeypatch.setattr(experts_cls, "_supports_current_device", lambda: True)
1894
+ config = make_dummy_moe_config(
1895
+ num_experts=256,
1896
+ num_local_experts=64 if use_ep else 256,
1897
+ experts_per_token=6,
1898
+ hidden_dim=4096,
1899
+ intermediate_size=2048,
1900
+ activation=MoEActivation.SWIGLUOAI,
1901
+ )
1902
+ config = replace(
1903
+ config,
1904
+ routing_method=RoutingMethodType.Renormalize,
1905
+ moe_parallel_config=replace(
1906
+ config.moe_parallel_config,
1907
+ tp_size=1 if use_ep else 4,
1908
+ ep_size=4 if use_ep else 1,
1909
+ use_ep=use_ep,
1910
+ ),
1911
+ )
1912
+ assert not config.moe_parallel_config.use_all2all_kernels
1913
+ supported, reason = experts_cls.is_supported_config(
1914
+ experts_cls,
1915
+ config,
1916
+ kMxfp4Static,
1917
+ kFp8StaticTensorSym if backend_name == "AITER_MXFP4_FP8" else None,
1918
+ FusedMoEActivationFormat.Standard,
1919
+ )
1920
+
1921
+ assert supported is not use_ep
1922
+ if use_ep:
1923
+ assert "parallel config" in reason
1924
+
1925
+
1926
# Every activation-quantizing OCP MX scheme must map to a `quant_dtype` that
1927
# `moe_kernel_quantize_input` actually dispatches on. Its final `else` returns
1928
# the activation untouched, so a name it does not know (e.g. "mxfp6" instead of
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py
+1
-1
@@ -321,7 +321,7 @@ class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
321
moe_parallel_config: FusedMoEParallelConfig,
322
) -> bool:
323
return (
324
- not moe_parallel_config.use_all2all_kernels
324
+ not moe_parallel_config.use_ep
325
and not moe_parallel_config.enable_eplb
326
and moe_parallel_config.dp_size <= 1
327
)
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
+1
-1
@@ -246,7 +246,7 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
246
moe_parallel_config: FusedMoEParallelConfig,
247
) -> bool:
248
return (
249
- not moe_parallel_config.use_all2all_kernels
249
+ not moe_parallel_config.use_ep
250
and not moe_parallel_config.enable_eplb
251
and moe_parallel_config.dp_size <= 1
252
)