iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add MiniMax-M3 modelopt nvfp4 support (#46756)

Browse Source 
Signed-off-by: Xin Li <xinli@nvidia.com>
Signed-off-by: jasonlizhengjian <jasonlizhengjian@gmail.com>
Co-authored-by: Xin Li <xinli@nvidia.com>
This commit is contained in:
Jason Li
2026-06-29 09:29:39 -07:00
committed by GitHub
co-authored by Xin Li
parent c8fb2963bd
commit debec6440b
4 changed files with 93 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/quantization/test_modelopt.py
+6
View File
@@ -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=[],
),
)
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+59 -9
View File
@@ -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),
vllm/model_executor/layers/quantization/modelopt.py
+27
View File
@@ -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
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+1
View File
@@ -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]