diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 91de16f79c..d9cfce2c21 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -171,6 +171,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): replace_parameter(layer, "w13_weight", w13_new, prefer_copy=is_weight_update) replace_parameter(layer, "w2_weight", w2_new, prefer_copy=is_weight_update) + # AITER backend requires weights to be marked as shuffled. + if self.unquantized_backend == UnquantizedMoeBackend.AITER: + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True + if not is_weight_update: # Setup moe kernel only on the first call. For the unquantized # method, moe_quant_config is either the constant diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 1c9237d3f6..30e9a6344b 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -28,6 +28,7 @@ from vllm.model_executor.layers.fused_moe.config import ( ) from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + Fp8MoeBackend, convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, make_fp8_moe_quant_config, @@ -766,6 +767,11 @@ class Fp8MoEMethod(FusedMoEMethodBase): replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale) replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale) + # AITER backend requires weights to be marked as shuffled. + if self.fp8_backend == Fp8MoeBackend.AITER: + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True + self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config: assert self.experts_cls is not None diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index d6fef0b3d3..d3f3033be7 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -350,6 +350,11 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase): self.w13_precision_config = w13_scale self.w2_precision_config = w2_scale + # AITER backend requires weights to be marked as shuffled. + if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True + if w13_bias is not None and w2_bias is not None: replace_parameter(layer, "w13_bias", w13_bias) replace_parameter(layer, "w2_bias", w2_bias) @@ -678,6 +683,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self.w13_precision_config = w13_scale self.w2_precision_config = w2_scale + # AITER backend requires weights to be marked as shuffled. + if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True + if w13_bias is not None and w2_bias is not None: replace_parameter(layer, "w13_bias", w13_bias) replace_parameter(layer, "w2_bias", w2_bias)