diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 7c217cd56..95998c34b 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -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}, " diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 76bab05e7..d02de06f2 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -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), diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index d58db611c..6cefdc8ee 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -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, ) diff --git a/python/sglang/srt/models/minimax_m3.py b/python/sglang/srt/models/minimax_m3.py index 3ccafeeb4..a32eace30 100644 --- a/python/sglang/srt/models/minimax_m3.py +++ b/python/sglang/srt/models/minimax_m3.py @@ -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): diff --git a/python/sglang/srt/models/minimax_m3_vl.py b/python/sglang/srt/models/minimax_m3_vl.py index c3d563962..c2a7beaf5 100644 --- a/python/sglang/srt/models/minimax_m3_vl.py +++ b/python/sglang/srt/models/minimax_m3_vl.py @@ -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): diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index 04bee16b5..246a26802 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -14,8 +14,10 @@ import torch.nn as nn from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.logits_processor import should_apply_lm_head_quant_method from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES +from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod from sglang.srt.layers.quantization.modelopt_quant import ( ModelOptFp4Config, ModelOptFp4LinearMethod, @@ -23,6 +25,7 @@ from sglang.srt.layers.quantization.modelopt_quant import ( ModelOptNvFp4A16LinearMethod, ) from sglang.srt.model_loader.loader import ModelOptModelLoader +from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM from sglang.srt.models.utils import WeightsMapper from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_cuda_ci @@ -518,6 +521,66 @@ class TestParseQuantHfConfig(CustomTestCase): class TestModelOptMixedPrecisionConfig(CustomTestCase): + def test_minimax_mixed_precision_resolves_runtime_names_and_mxfp8(self): + quant_config = ModelOptMixedPrecisionConfig.from_config( + { + "quant_algo": "MIXED_PRECISION", + "weight_block_size": [1, 32], + "exclude_modules": ["language_model.lm_head"], + "quantized_layers": { + "language_model.model.layers.3.self_attn.q_proj": { + "quant_algo": "MXFP8" + }, + "language_model.model.layers.3.self_attn.k_proj": { + "quant_algo": "MXFP8" + }, + "language_model.model.layers.3.self_attn.v_proj": { + "quant_algo": "MXFP8" + }, + "language_model.model.layers.3.block_sparse_moe.experts.0.w1": { + "quant_algo": "NVFP4", + "group_size": 16, + }, + "language_model.model.layers.3.block_sparse_moe.shared_experts.gate_proj": { + "quant_algo": "MXFP8" + }, + }, + "packed_modules_mapping": { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + }, + } + ) + quant_config.apply_weight_name_mapper( + MiniMaxM3SparseForCausalLM.hf_to_sglang_mapper + ) + + self.assertEqual( + quant_config._resolve_quant_algo( + "language_model.model.layers.3.mlp.experts" + ), + "NVFP4", + ) + self.assertEqual( + quant_config._resolve_quant_algo( + "language_model.model.layers.3.mlp.shared_experts.gate_up_proj" + ), + "MXFP8", + ) + + # Type dispatch only needs a LinearBase instance; skip GPU weight setup. + linear = ReplicatedLinear.__new__(ReplicatedLinear) + method = quant_config.get_quant_method( + linear, "language_model.model.layers.3.self_attn.qkv_proj" + ) + self.assertIsInstance(method, Fp8LinearMethod) + self.assertTrue(method.use_mxfp8) + self.assertEqual(quant_config.mxfp8_config.weight_block_size, [1, 32]) + self.assertEqual( + quant_config.exclude_modules, + ["language_model.lm_head", "lm_head"], + ) + def test_nemotron_mixed_precision_with_nvfp4_layers_uses_modelopt_mixed(self): model_config = ModelConfig.__new__(ModelConfig) model_config.hf_config = MagicMock()