A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 07:45 UTC
[ROCm][Bugfix] Fall back from CK MXFP4 MoE when GEMM dimensions are unsupported (#35893)
Signed-off-by: Li <chuali@amd.com>
Chuan (Richard) Li committed
Mar 4, 2026 at 00:30 UTC
5dc3538736e40428d388f6980025340f4f7524af
3 files changed
+63
-1
vllm/model_executor/layers/quantization/mxfp4.py
+26
@@ -48,6 +48,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
48
prepare_moe_fp4_layer_for_marlin,
49
)
50
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
51
+ CK_MXFP4_MOE_DIM_ALIGNMENT,
52
_can_support_mxfp4,
53
_swizzle_mxfp4,
54
get_padding_alignment,
@@ -259,6 +260,31 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
260
get_current_vllm_config().compilation_config.max_cudagraph_capture_size
261
)
262
263
+ # CK's pre-compiled MXFP4 MoE GEMM kernel instances have dimension
264
+ # alignment requirements. Fall back to Triton when not met.
265
+ if (
266
+ self.mxfp4_backend == Mxfp4Backend.CK
267
+ and moe.intermediate_size_per_partition % CK_MXFP4_MOE_DIM_ALIGNMENT != 0
268
+ ):
269
+ if has_triton_kernels():
270
+ logger.warning_once(
271
+ "CK MXFP4 MoE GEMM does not support "
272
+ "intermediate_size_per_partition=%d (not a multiple of "
273
+ "%d). Falling back to Triton backend.",
274
+ moe.intermediate_size_per_partition,
275
+ CK_MXFP4_MOE_DIM_ALIGNMENT,
276
+ )
277
+ self.mxfp4_backend = Mxfp4Backend.TRITON
278
+ else:
279
+ raise ValueError(
280
+ f"CK MXFP4 MoE GEMM does not support "
281
+ f"intermediate_size_per_partition="
282
+ f"{moe.intermediate_size_per_partition} (not a multiple "
283
+ f"of {CK_MXFP4_MOE_DIM_ALIGNMENT}) and no Triton "
284
+ f"fallback is available. Use a compatible "
285
+ f"tensor_parallel_size."
286
+ )
287
+
288
assert self.mxfp4_backend != Mxfp4Backend.NONE, (
289
f"get_mxfp4_backend(with_lora_support={moe.is_lora_enabled}) found"
290
"no compatible MXFP4 MoE backend (FlashInfer/Marlin/Triton)."
vllm/model_executor/layers/quantization/quark/quark_moe.py
+30
-1
@@ -32,7 +32,10 @@ from vllm.model_executor.layers.quantization.mxfp4 import (
32
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
33
prepare_fp8_moe_layer_for_marlin,
34
)
35
-from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4
35
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
36
+ CK_MXFP4_MOE_DIM_ALIGNMENT,
37
+ _swizzle_mxfp4,
38
+)
39
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
40
OCP_MX_BLOCK_SIZE,
41
OCP_MX_Scheme,
@@ -732,6 +735,32 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
735
or not self.ocp_mx_scheme.startswith("w_mxfp4")
736
) and (self.mxfp4_backend is None or not self.use_rocm_aiter_moe)
737
738
+ # CK's pre-compiled MXFP4 MoE GEMM kernel instances have dimension
739
+ # alignment requirements. When violated (e.g. MiniMax-M2.1 with
740
+ # TP=4 yields intermediate_size_per_partition=384), AITER raises:
741
+ # "device_gemm ... does not support this GEMM problem".
742
+ # Fall back to emulation in that case.
743
+ if (
744
+ not self.emulate
745
+ and self.use_rocm_aiter_moe
746
+ and self.ocp_mx_scheme is not None
747
+ and self.ocp_mx_scheme.startswith("w_mxfp4")
748
+ and moe.intermediate_size_per_partition % CK_MXFP4_MOE_DIM_ALIGNMENT != 0
749
+ ):
750
+ logger.warning_once(
751
+ "AITER CK MXFP4 MoE GEMM does not support "
752
+ "intermediate_size_per_partition=%d (not a multiple of %d). "
753
+ "This typically happens when intermediate_size / "
754
+ "tensor_parallel_size produces an incompatible dimension. "
755
+ "Falling back to emulation mode. To avoid this overhead, "
756
+ "use a compatible tensor_parallel_size or set "
757
+ "VLLM_ROCM_USE_AITER_MOE=0.",
758
+ moe.intermediate_size_per_partition,
759
+ CK_MXFP4_MOE_DIM_ALIGNMENT,
760
+ )
761
+ self.use_rocm_aiter_moe = False
762
+ self.emulate = True
763
+
764
if self.emulate:
765
logger.warning_once(
766
f"The current mode (supports_mx={current_platform.supports_mx()}, "
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+7
@@ -14,6 +14,13 @@ from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_
14
15
logger = init_logger(__name__)
16
17
+# CK's pre-compiled MXFP4 MoE GEMM kernel instances require the
18
+# intermediate_size (after TP split) to be a multiple of this value.
19
+# This arises from FP4 packing (2 values per byte) combined with CK
20
+# tile size constraints. When violated, AITER raises:
21
+# "device_gemm ... does not support this GEMM problem".
22
+CK_MXFP4_MOE_DIM_ALIGNMENT = 256
23
+
24
25
def _swizzle_mxfp4(quant_tensor, scale, num_warps):
26
"""weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""