minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat: Support nvidia/MiniMax-M3-NVFP4 (#31989)

Browse Source
This commit is contained in:
Trevor Morris
2026-07-30 14:32:03 -07:00
committed by GitHub
parent c4af6cf263
commit a6221d776f
6 changed files with 156 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/overrides.py
+5
View File
@@ -528,6 +528,11 @@ def _minimax_m3_overrides(server_args: Any, hf_config: Any) -> dict:
page_resolved = 128
if server_args.moe_runner_backend == "auto" and quant_resolved == "mxfp8":
overrides["moe_runner_backend"] = "deep_gemm"
elif (
server_args.moe_runner_backend == "auto"
and quant_resolved == "modelopt_mixed"
):
overrides["moe_runner_backend"] = "flashinfer_trtllm_routed"
logger.info(
"MiniMax-M3 on SM100: attention_backend="
f"{overrides.get('attention_backend', server_args.attention_backend)}, page_size={page_resolved}, "
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+6 -4
View File
@@ -886,6 +886,8 @@ class FlashInferTrtllmFp4MoeQuantInfo(MoeQuantInfo):
routing_method_type: int
use_per_token_activation: bool = False
gemm1_alpha: Optional[torch.Tensor] = None
gemm1_beta: Optional[torch.Tensor] = None
gemm1_clamp_limit: Optional[torch.Tensor] = None
@@ -1053,8 +1055,8 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
gemm1_weights=quant_info.w13_weight,
gemm1_weights_scale=quant_info.w13_weight_scale.view(torch.float8_e4m3fn),
gemm1_bias=None,
gemm1_alpha=None,
gemm1_beta=None,
gemm1_alpha=quant_info.gemm1_alpha,
gemm1_beta=quant_info.gemm1_beta,
gemm1_clamp_limit=quant_info.gemm1_clamp_limit,
gemm2_weights=quant_info.w2_weight,
gemm2_weights_scale=quant_info.w2_weight_scale.view(torch.float8_e4m3fn),
@@ -1094,8 +1096,8 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
gemm1_weights=quant_info.w13_weight,
gemm1_weights_scale=quant_info.w13_weight_scale.view(torch.float8_e4m3fn),
gemm1_bias=None,
gemm1_alpha=None,
gemm1_beta=None,
gemm1_alpha=quant_info.gemm1_alpha,
gemm1_beta=quant_info.gemm1_beta,
gemm1_clamp_limit=quant_info.gemm1_clamp_limit,
gemm2_weights=quant_info.w2_weight,
gemm2_weights_scale=quant_info.w2_weight_scale.view(torch.float8_e4m3fn),
python/sglang/srt/layers/quantization/modelopt_quant.py
+46 -12
View File
@@ -34,7 +34,7 @@ from sglang.srt.layers.quantization.fp4_utils import (
fp4_quantize,
get_fp4_gemm_runner_backend,
)
from sglang.srt.layers.quantization.fp8 import Fp8Config
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod, Fp8MoEMethod
from sglang.srt.layers.quantization.fp8_utils import (
apply_fp8_linear,
apply_fp8_linear_bmm_flashinfer,
@@ -634,10 +634,12 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
fp8_config: ModelOptFp8Config,
nvfp4_config: ModelOptFp4Config,
nvfp4a16_config: ModelOptFp4Config,
mxfp8_config: Fp8Config,
) -> None:
super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping)
self.quantized_layers = quantized_layers
self.fp8_config = fp8_config
self.mxfp8_config = mxfp8_config
self.nvfp4_config = nvfp4_config
self.nvfp4a16_config = nvfp4a16_config
@@ -685,7 +687,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
kv_cache_quant_algo = "auto"
else:
kv_cache_quant_algo = config.get("kv_cache_quant_algo")
exclude_modules = config.get("ignore")
exclude_modules = config.get("ignore", config.get("exclude_modules"))
quantized_layers = config.get("quantized_layers", {})
else:
quantization_section = cls.get_from_keys(config, ["quantization"])
@@ -721,6 +723,13 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
exclude_modules=[],
packed_modules_mapping=packed_modules_mapping,
)
mxfp8_config = Fp8Config(
is_checkpoint_fp8_serialized=True,
activation_scheme="dynamic",
weight_block_size=[1, 32],
packed_modules_mapping=packed_modules_mapping,
use_mxfp8=True,
)
nvfp4_config = ModelOptFp4Config(
is_checkpoint_nvfp4_serialized=True,
kv_cache_quant_algo=kv_cache_quant_algo,
@@ -743,6 +752,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
packed_modules_mapping=packed_modules_mapping,
quantized_layers=quantized_layers,
fp8_config=fp8_config,
mxfp8_config=mxfp8_config,
nvfp4_config=nvfp4_config,
nvfp4a16_config=nvfp4a16_config,
)
@@ -820,6 +830,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
return UnquantizedLinearMethod()
if quant_algo == "FP8":
return ModelOptFp8LinearMethod(self.fp8_config)
if quant_algo == "MXFP8":
return Fp8LinearMethod(self.mxfp8_config)
if quant_algo == "NVFP4":
return ModelOptFp4LinearMethod(self.nvfp4_config)
if quant_algo == "W4A16_NVFP4":
@@ -834,6 +846,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
return None
if quant_algo == "FP8":
return ModelOptFp8MoEMethod(self.fp8_config)
if quant_algo == "MXFP8":
return Fp8MoEMethod(self.mxfp8_config)
if quant_algo == "NVFP4":
return ModelOptNvFp4FusedMoEMethod(self.nvfp4_config)
if quant_algo == "W4A16_NVFP4":
@@ -2282,17 +2296,33 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
(1 / w2_input_scale).to(torch.float32),
)
swiglu_limit = layer.moe_runner_config.swiglu_limit
if (
swiglu_limit is not None
and layer.moe_runner_config.is_gated
and self.enable_flashinfer_trtllm_moe
):
copy_or_rebind_param(
layer,
"gemm1_clamp_limit",
(swiglu_limit / layer.g1_alphas).to(torch.float32),
if layer.moe_runner_config.is_gated and self.enable_flashinfer_trtllm_moe:
gemm1_clamp_limit = (
layer.moe_runner_config.gemm1_clamp_limit
or layer.moe_runner_config.swiglu_limit
)
if gemm1_clamp_limit is not None:
copy_or_rebind_param(
layer,
"gemm1_clamp_limit",
(gemm1_clamp_limit / layer.g1_alphas).to(torch.float32),
)
if layer.moe_runner_config.gemm1_alpha is not None:
copy_or_rebind_param(
layer,
"gemm1_alpha",
torch.full_like(
layer.g1_alphas,
layer.moe_runner_config.gemm1_alpha,
dtype=torch.float32,
),
)
copy_or_rebind_param(
layer,
"gemm1_beta",
(1.0 / layer.g1_alphas).to(torch.float32),
)
# TODO: for flashinfer always do MOE_NVFP4_DISPATCH
layer.dispatcher.set_quant_config(
@@ -2570,6 +2600,8 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
)
gemm1_clamp = getattr(layer, "gemm1_clamp_limit", None)
gemm1_alpha = getattr(layer, "gemm1_alpha", None)
gemm1_beta = getattr(layer, "gemm1_beta", None)
quant_info = FlashInferTrtllmFp4MoeQuantInfo(
w13_weight=layer.w13_weight.data,
w2_weight=layer.w2_weight.data,
@@ -2585,6 +2617,8 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
intermediate_size_per_partition=layer.intermediate_size_per_partition,
routing_method_type=routing_method_type,
use_per_token_activation=self.quant_config.use_per_token_activation,
gemm1_alpha=gemm1_alpha.data if gemm1_alpha is not None else None,
gemm1_beta=gemm1_beta.data if gemm1_beta is not None else None,
gemm1_clamp_limit=gemm1_clamp.data if gemm1_clamp is not None else None,
)
python/sglang/srt/models/minimax_m3.py
+18
View File
@@ -79,6 +79,7 @@ from sglang.srt.model_loader.weight_utils import (
maybe_remap_kv_scale_name,
)
from sglang.srt.models.minimax_m2 import MiniMaxM2RMSNormTP
from sglang.srt.models.utils import WeightsMapper
from sglang.srt.runtime_context import get_parallel, get_server_args
from sglang.srt.utils import (
add_prefix,
@@ -1418,6 +1419,15 @@ class MiniMaxM3Model(nn.Module):
class MiniMaxM3SparseForCausalLM(nn.Module):
hf_to_sglang_mapper = WeightsMapper(
orig_to_new_substr={".block_sparse_moe.": ".mlp."}
)
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"index_qkv_proj": ["index_q_proj", "index_k_proj", "index_v_proj"],
"gate_up_proj": ["gate_proj", "up_proj"],
}
def __init__(
self,
config: PretrainedConfig,
@@ -1462,6 +1472,14 @@ class MiniMaxM3SparseForCausalLM(nn.Module):
disable_reason = None
if not getattr(self.config, "n_shared_experts", None):
disable_reason = "No shared experts are defined in the config."
elif (
self.quant_config is not None
and self.quant_config.get_name() == "modelopt_mixed"
):
disable_reason = (
"Shared and routed experts may use different quantization formats "
"in ModelOpt mixed-precision checkpoints."
)
elif not _is_cuda:
disable_reason = "Shared experts fusion currently requires CUDA devices."
elif _is_cuda and (_device_sm is not None) and (_device_sm < 80):
python/sglang/srt/models/minimax_m3_vl.py
+18
View File
@@ -42,6 +42,7 @@ from sglang.srt.models.minimax_vl_common import (
load_vision_weight,
merge_vit_qkv_weights,
)
from sglang.srt.models.utils import WeightsMapper
from sglang.srt.runtime_context import get_parallel, get_server_args
from sglang.srt.utils import add_prefix, get_device_sm, is_cuda, log_info_on_rank0
from sglang.srt.utils.hf_transformers_utils import get_rope_config
@@ -54,6 +55,15 @@ _device_sm = get_device_sm()
class MiniMaxM3SparseForConditionalGeneration(nn.Module):
hf_to_sglang_mapper = WeightsMapper(
orig_to_new_substr={".block_sparse_moe.": ".mlp."}
)
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"index_qkv_proj": ["index_q_proj", "index_k_proj", "index_v_proj"],
"gate_up_proj": ["gate_proj", "up_proj"],
}
def __init__(
self,
config,
@@ -131,6 +141,14 @@ class MiniMaxM3SparseForConditionalGeneration(nn.Module):
disable_reason = None
if not getattr(text_config, "n_shared_experts", None):
disable_reason = "No shared experts are defined in the config."
elif (
self.quant_config is not None
and self.quant_config.get_name() == "modelopt_mixed"
):
disable_reason = (
"Shared and routed experts may use different quantization formats "
"in ModelOpt mixed-precision checkpoints."
)
elif not _is_cuda:
disable_reason = "Shared experts fusion currently requires CUDA devices."
elif (_device_sm is not None) and (_device_sm < 80):