From 944e138bcf39e9236bbfd49d98f00fb45e6cea54 Mon Sep 17 00:00:00 2001 From: Rohan Potdar <66227218+Rohan138@users.noreply.github.com> Date: Wed, 29 Apr 2026 16:39:03 -0500 Subject: [PATCH] [ROCm][Bugfix]: W4A4 MOE using emulation instead of AITER on MXFP4-supported hardware (#41175) Signed-off-by: Rohan138 --- .../layers/fused_moe/rocm_aiter_fused_moe.py | 3 ++- .../model_executor/layers/quantization/quark/quark_moe.py | 8 ++------ 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py index 495b9daaff4..d9d888296b7 100644 --- a/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py @@ -252,7 +252,8 @@ def rocm_aiter_fused_experts( else: quant_method = QuantMethod.NO.value - # mxfp4: both w4a4 (quark) and w4a16 (oracle CK) use BLOCK_1X32 + # mxfp4 i.e. w4a4, w4a16 uses BLOCK_1X32 + # mxfp6 and mxfp8 are unsupported in AITER currently and use emulation instead if quant_config.use_mxfp4_w4a4 or quant_config.use_mxfp4_w4a16: quant_method = QuantMethod.BLOCK_1X32.value # w8a8 block-scaled diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index d92acb85c26..8737d9e7372 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1025,12 +1025,8 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): get_current_vllm_config().model_config.hf_config, "model_type", None ) - # 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",) + # TODO: Remove once all OCP MX schemes use the kernel abstraction + _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