Add MiniMax-M3 modelopt nvfp4 support (#46756)
Signed-off-by: Xin Li <xinli@nvidia.com> Signed-off-by: jasonlizhengjian <jasonlizhengjian@gmail.com> Co-authored-by: Xin Li <xinli@nvidia.com>
Jason Li committed
Jun 29, 2026 at 12:29 UTC
debec6440b89fe6ab14acb00e6eb2b04257f57a2
4 files changed
+93
-9
tests/quantization/test_modelopt.py
+6
@@ -18,6 +18,7 @@ from vllm.model_executor.layers.linear import UnquantizedLinearMethod
18
from vllm.model_executor.layers.quantization.modelopt import (
19
ModelOptFp8Config,
20
ModelOptMixedPrecisionConfig,
21
+ ModelOptMxFp8Config,
22
ModelOptNvFp4Config,
23
ModelOptNvFp4LinearMethod,
24
)
@@ -84,6 +85,11 @@ def _mixed_precision_config(quantized_layers: dict) -> ModelOptMixedPrecisionCon
85
kv_cache_quant_algo=None,
86
exclude_modules=[],
87
),
88
+ mxfp8_config=ModelOptMxFp8Config(
89
+ is_checkpoint_mxfp8_serialized=True,
90
+ kv_cache_quant_algo=None,
91
+ exclude_modules=[],
92
+ ),
93
)
94
95
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+59
-9
@@ -66,16 +66,46 @@ class TrtLlmNvFp4ExpertsBase:
66
else:
67
self.g1_scale_c = self.quant_config.a2_gscale.clone()
68
69
- if moe_config.is_act_and_mul and quant_config.gemm1_clamp_limit is not None:
70
- device = torch.accelerator.current_device_index()
71
- self.gemm1_clamp_limit = torch.full(
69
+ # Fall back to moe_config.swiglu_* when quant_config doesn't carry them
70
+ # (ModelOpt NVFP4 checkpoints store these on moe_config, not quant_config).
71
+ device = torch.accelerator.current_device_index()
72
+
73
+ def _per_expert(val: float | None) -> torch.Tensor | None:
74
+ if val is None:
75
+ return None
76
+ return torch.full(
77
(self.local_num_experts,),
73
- quant_config.gemm1_clamp_limit,
78
+ float(val),
79
dtype=torch.float32,
80
device=device,
81
)
82
+
83
+ clamp = quant_config.gemm1_clamp_limit
84
+ if clamp is None:
85
+ clamp = getattr(moe_config, "swiglu_limit", None)
86
+ alpha = quant_config.gemm1_alpha
87
+ if alpha is None:
88
+ alpha = getattr(moe_config, "swiglu_alpha", None)
89
+ beta = quant_config.gemm1_beta
90
+ if beta is None:
91
+ beta = getattr(moe_config, "swiglu_beta", None)
92
+
93
+ if moe_config.is_act_and_mul:
94
+ self.gemm1_clamp_limit = _per_expert(clamp)
95
+ self.gemm1_alpha = _per_expert(alpha)
96
+ self.gemm1_beta = _per_expert(beta)
97
else:
98
self.gemm1_clamp_limit = None
99
+ self.gemm1_alpha = None
100
+ self.gemm1_beta = None
101
+
102
+ logger.debug_once(
103
+ "activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s",
104
+ moe_config.activation,
105
+ alpha,
106
+ beta,
107
+ clamp,
108
+ )
109
110
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
111
layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
@@ -109,6 +139,25 @@ class TrtLlmNvFp4ExpertsBase:
139
)
140
self.gemm1_clamp_limit = layer.gemm1_clamp_limit
141
142
+ # beta shifts the raw GEMM1 accumulator, so fold by g1_alphas like the
143
+ # clamp limit. alpha is applied to the dequantized gate, so it stays
144
+ # raw. Register both on the layer so EPLB rearranges them with the
145
+ # other per-expert tensors.
146
+ if self.gemm1_beta is not None:
147
+ gemm1_beta = self.gemm1_beta / self.quant_config.g1_alphas
148
+ layer.register_parameter(
149
+ "gemm1_beta",
150
+ torch.nn.Parameter(gemm1_beta, requires_grad=False),
151
+ )
152
+ self.gemm1_beta = layer.gemm1_beta
153
+
154
+ if self.gemm1_alpha is not None:
155
+ layer.register_parameter(
156
+ "gemm1_alpha",
157
+ torch.nn.Parameter(self.gemm1_alpha, requires_grad=False),
158
+ )
159
+ self.gemm1_alpha = layer.gemm1_alpha
160
+
161
@staticmethod
162
def _supports_current_device() -> bool:
163
"""Supports only Blackwell-family GPUs."""
@@ -137,12 +186,13 @@ class TrtLlmNvFp4ExpertsBase:
186
187
@staticmethod
188
def _supports_activation(activation: MoEActivation) -> bool:
140
- """Supports only SiLU, RELU^2 non-gated and GELU activation."""
189
+ """Supports SiLU, RELU^2 non-gated, GELU, and clamped SwiGLU-OAI."""
190
return activation in [
191
MoEActivation.SILU,
192
MoEActivation.RELU2_NO_MUL,
193
MoEActivation.GELU,
194
MoEActivation.GELU_TANH,
195
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
196
]
197
198
@staticmethod
@@ -248,8 +298,8 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
298
gemm1_weights=w1,
299
gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn),
300
gemm1_bias=None,
251
- gemm1_alpha=None,
252
- gemm1_beta=None,
301
+ gemm1_alpha=self.gemm1_alpha,
302
+ gemm1_beta=self.gemm1_beta,
303
gemm1_clamp_limit=self.gemm1_clamp_limit,
304
gemm2_weights=w2,
305
gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn),
@@ -409,8 +459,8 @@ class TrtLlmNvFp4ExpertsMonolithic(
459
gemm1_weights=w1,
460
gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn),
461
gemm1_bias=None,
412
- gemm1_alpha=None,
413
- gemm1_beta=None,
462
+ gemm1_alpha=self.gemm1_alpha,
463
+ gemm1_beta=self.gemm1_beta,
464
gemm1_clamp_limit=self.gemm1_clamp_limit,
465
gemm2_weights=w2,
466
gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn),
vllm/model_executor/layers/quantization/modelopt.py
+27
@@ -2283,6 +2283,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2283
fp8_config: ModelOptFp8Config,
2284
nvfp4_config: ModelOptNvFp4Config,
2285
w4a16_nvfp4_config: ModelOptNvFp4Config,
2286
+ mxfp8_config: ModelOptMxFp8Config,
2287
) -> None:
2288
super().__init__(exclude_modules)
2289
self.kv_cache_quant_method = kv_cache_quant_method
@@ -2290,6 +2291,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2291
self.fp8_config = fp8_config
2292
self.nvfp4_config = nvfp4_config
2293
self.w4a16_nvfp4_config = w4a16_nvfp4_config
2294
+ self.mxfp8_config = mxfp8_config
2295
2296
def get_name(self) -> QuantizationMethods:
2297
return "modelopt_mixed"
@@ -2380,6 +2382,12 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2382
group_size=group_size,
2383
)
2384
2385
+ mxfp8_config = ModelOptMxFp8Config(
2386
+ is_checkpoint_mxfp8_serialized=True,
2387
+ kv_cache_quant_algo=kv_cache_quant_method,
2388
+ exclude_modules=[],
2389
+ )
2390
+
2391
return cls(
2392
kv_cache_quant_method=kv_cache_quant_method,
2393
exclude_modules=exclude_modules,
@@ -2387,6 +2395,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2395
fp8_config=fp8_config,
2396
nvfp4_config=nvfp4_config,
2397
w4a16_nvfp4_config=w4a16_nvfp4_config,
2398
+ mxfp8_config=mxfp8_config,
2399
)
2400
2401
def _resolve_quant_algo(self, prefix: str) -> str | None:
@@ -2441,6 +2450,17 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2450
if key.startswith(parent_dot):
2451
return info["quant_algo"].upper()
2452
2453
+ # 4. Parent-prefix fallback for fused projections (qkv_proj, gate_up_proj).
2454
+ for candidate in self._quantized_layer_prefix_candidates(prefix):
2455
+ parent_dot = candidate.rsplit(".", 1)[0] + "."
2456
+ algos = {
2457
+ info["quant_algo"].upper()
2458
+ for key, info in self.quantized_layers.items()
2459
+ if key.startswith(parent_dot) and "." not in key[len(parent_dot) :]
2460
+ }
2461
+ if len(algos) == 1:
2462
+ return algos.pop()
2463
+
2464
return None
2465
2466
@staticmethod
@@ -2486,6 +2506,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2506
return ModelOptNvFp4LinearMethod(self.nvfp4_config)
2507
if quant_algo == "W4A16_NVFP4":
2508
return ModelOptNvFp4W4A16LinearMethod(self.w4a16_nvfp4_config)
2509
+ if quant_algo == "MXFP8":
2510
+ return ModelOptMxFp8LinearMethod(self.mxfp8_config)
2511
# Layer not in quantized_layers — leave unquantized
2512
return UnquantizedLinearMethod()
2513
@@ -2505,6 +2527,11 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2527
quant_config=self.w4a16_nvfp4_config,
2528
moe_config=layer.moe_config,
2529
)
2530
+ if quant_algo == "MXFP8":
2531
+ return ModelOptMxFp8FusedMoE(
2532
+ quant_config=self.mxfp8_config,
2533
+ moe_config=layer.moe_config,
2534
+ )
2535
return None
2536
2537
return None
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+1
@@ -36,6 +36,7 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType"
36
MoEActivation.GELU: ActivationType.Geglu,
37
MoEActivation.GELU_TANH: ActivationType.Geglu,
38
MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
39
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu,
40
}
41
return ACTIVATION_TO_FI_ACTIVATION[activation]
42