From debec6440b89fe6ab14acb00e6eb2b04257f57a2 Mon Sep 17 00:00:00 2001 From: Jason Li Date: Mon, 29 Jun 2026 12:29:39 -0400 Subject: [PATCH] Add MiniMax-M3 modelopt nvfp4 support (#46756) Signed-off-by: Xin Li Signed-off-by: jasonlizhengjian Co-authored-by: Xin Li --- tests/quantization/test_modelopt.py | 6 ++ .../fused_moe/experts/trtllm_nvfp4_moe.py | 68 ++++++++++++++++--- .../layers/quantization/modelopt.py | 27 ++++++++ .../quantization/utils/flashinfer_utils.py | 1 + 4 files changed, 93 insertions(+), 9 deletions(-) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 0b54bcdbdf..3245023148 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -18,6 +18,7 @@ from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.modelopt import ( ModelOptFp8Config, ModelOptMixedPrecisionConfig, + ModelOptMxFp8Config, ModelOptNvFp4Config, ModelOptNvFp4LinearMethod, ) @@ -84,6 +85,11 @@ def _mixed_precision_config(quantized_layers: dict) -> ModelOptMixedPrecisionCon kv_cache_quant_algo=None, exclude_modules=[], ), + mxfp8_config=ModelOptMxFp8Config( + is_checkpoint_mxfp8_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=[], + ), ) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index e45fc77ad9..518c87ce4d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -66,16 +66,46 @@ class TrtLlmNvFp4ExpertsBase: else: self.g1_scale_c = self.quant_config.a2_gscale.clone() - if moe_config.is_act_and_mul and quant_config.gemm1_clamp_limit is not None: - device = torch.accelerator.current_device_index() - self.gemm1_clamp_limit = torch.full( + # Fall back to moe_config.swiglu_* when quant_config doesn't carry them + # (ModelOpt NVFP4 checkpoints store these on moe_config, not quant_config). + device = torch.accelerator.current_device_index() + + def _per_expert(val: float | None) -> torch.Tensor | None: + if val is None: + return None + return torch.full( (self.local_num_experts,), - quant_config.gemm1_clamp_limit, + float(val), dtype=torch.float32, device=device, ) + + clamp = quant_config.gemm1_clamp_limit + if clamp is None: + clamp = getattr(moe_config, "swiglu_limit", None) + alpha = quant_config.gemm1_alpha + if alpha is None: + alpha = getattr(moe_config, "swiglu_alpha", None) + beta = quant_config.gemm1_beta + if beta is None: + beta = getattr(moe_config, "swiglu_beta", None) + + if moe_config.is_act_and_mul: + self.gemm1_clamp_limit = _per_expert(clamp) + self.gemm1_alpha = _per_expert(alpha) + self.gemm1_beta = _per_expert(beta) else: self.gemm1_clamp_limit = None + self.gemm1_alpha = None + self.gemm1_beta = None + + logger.debug_once( + "activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s", + moe_config.activation, + alpha, + beta, + clamp, + ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) @@ -109,6 +139,25 @@ class TrtLlmNvFp4ExpertsBase: ) self.gemm1_clamp_limit = layer.gemm1_clamp_limit + # beta shifts the raw GEMM1 accumulator, so fold by g1_alphas like the + # clamp limit. alpha is applied to the dequantized gate, so it stays + # raw. Register both on the layer so EPLB rearranges them with the + # other per-expert tensors. + if self.gemm1_beta is not None: + gemm1_beta = self.gemm1_beta / self.quant_config.g1_alphas + layer.register_parameter( + "gemm1_beta", + torch.nn.Parameter(gemm1_beta, requires_grad=False), + ) + self.gemm1_beta = layer.gemm1_beta + + if self.gemm1_alpha is not None: + layer.register_parameter( + "gemm1_alpha", + torch.nn.Parameter(self.gemm1_alpha, requires_grad=False), + ) + self.gemm1_alpha = layer.gemm1_alpha + @staticmethod def _supports_current_device() -> bool: """Supports only Blackwell-family GPUs.""" @@ -137,12 +186,13 @@ class TrtLlmNvFp4ExpertsBase: @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - """Supports only SiLU, RELU^2 non-gated and GELU activation.""" + """Supports SiLU, RELU^2 non-gated, GELU, and clamped SwiGLU-OAI.""" return activation in [ MoEActivation.SILU, MoEActivation.RELU2_NO_MUL, MoEActivation.GELU, MoEActivation.GELU_TANH, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, ] @staticmethod @@ -248,8 +298,8 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula gemm1_weights=w1, gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, gemm1_clamp_limit=self.gemm1_clamp_limit, gemm2_weights=w2, gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), @@ -409,8 +459,8 @@ class TrtLlmNvFp4ExpertsMonolithic( gemm1_weights=w1, gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn), gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, gemm1_clamp_limit=self.gemm1_clamp_limit, gemm2_weights=w2, gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn), diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index d51a2dd312..8fa1cb4d54 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -2283,6 +2283,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): fp8_config: ModelOptFp8Config, nvfp4_config: ModelOptNvFp4Config, w4a16_nvfp4_config: ModelOptNvFp4Config, + mxfp8_config: ModelOptMxFp8Config, ) -> None: super().__init__(exclude_modules) self.kv_cache_quant_method = kv_cache_quant_method @@ -2290,6 +2291,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): self.fp8_config = fp8_config self.nvfp4_config = nvfp4_config self.w4a16_nvfp4_config = w4a16_nvfp4_config + self.mxfp8_config = mxfp8_config def get_name(self) -> QuantizationMethods: return "modelopt_mixed" @@ -2380,6 +2382,12 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): group_size=group_size, ) + mxfp8_config = ModelOptMxFp8Config( + is_checkpoint_mxfp8_serialized=True, + kv_cache_quant_algo=kv_cache_quant_method, + exclude_modules=[], + ) + return cls( kv_cache_quant_method=kv_cache_quant_method, exclude_modules=exclude_modules, @@ -2387,6 +2395,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): fp8_config=fp8_config, nvfp4_config=nvfp4_config, w4a16_nvfp4_config=w4a16_nvfp4_config, + mxfp8_config=mxfp8_config, ) def _resolve_quant_algo(self, prefix: str) -> str | None: @@ -2441,6 +2450,17 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): if key.startswith(parent_dot): return info["quant_algo"].upper() + # 4. Parent-prefix fallback for fused projections (qkv_proj, gate_up_proj). + for candidate in self._quantized_layer_prefix_candidates(prefix): + parent_dot = candidate.rsplit(".", 1)[0] + "." + algos = { + info["quant_algo"].upper() + for key, info in self.quantized_layers.items() + if key.startswith(parent_dot) and "." not in key[len(parent_dot) :] + } + if len(algos) == 1: + return algos.pop() + return None @staticmethod @@ -2486,6 +2506,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): return ModelOptNvFp4LinearMethod(self.nvfp4_config) if quant_algo == "W4A16_NVFP4": return ModelOptNvFp4W4A16LinearMethod(self.w4a16_nvfp4_config) + if quant_algo == "MXFP8": + return ModelOptMxFp8LinearMethod(self.mxfp8_config) # Layer not in quantized_layers — leave unquantized return UnquantizedLinearMethod() @@ -2505,6 +2527,11 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): quant_config=self.w4a16_nvfp4_config, moe_config=layer.moe_config, ) + if quant_algo == "MXFP8": + return ModelOptMxFp8FusedMoE( + quant_config=self.mxfp8_config, + moe_config=layer.moe_config, + ) return None return None diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 1cbfdf69c9..9961d0f0a1 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -36,6 +36,7 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType" MoEActivation.GELU: ActivationType.Geglu, MoEActivation.GELU_TANH: ActivationType.Geglu, MoEActivation.RELU2_NO_MUL: ActivationType.Relu2, + MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu, } return ACTIVATION_TO_FI_ACTIVATION[activation]