From 75a7879fd4c149e675b50f9b21964e557577cc4f Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Sat, 14 Mar 2026 15:56:26 +0800 Subject: [PATCH] [Model] Support Nemotron 3 Super NVFP4 (#20407) --- python/sglang/srt/configs/model_config.py | 12 ++ .../srt/layers/quantization/__init__.py | 2 + .../srt/layers/quantization/modelopt_quant.py | 177 ++++++++++++++++++ python/sglang/srt/model_loader/loader.py | 6 +- python/sglang/srt/server_args.py | 26 ++- .../model_loading/test_modelopt_loader.py | 65 +++++++ 6 files changed, 277 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index c5389ee69..099370c3e 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -793,6 +793,11 @@ class ModelConfig: quant_algo = json_quant_configs.get("quant_algo", None) if quant_algo == "MIXED_PRECISION": + architectures = getattr(self.hf_config, "architectures", []) or [] + if getattr(self.hf_config, "model_type", None) == "nemotron_h" or any( + arch.startswith("NemotronH") for arch in architectures + ): + return {"quant_method": "modelopt_mixed", "quant_algo": quant_algo} return {"quant_method": "w4afp8", "quant_algo": quant_algo} elif quant_algo and ("FP4" in quant_algo or "NVFP4" in quant_algo): return {"quant_method": "modelopt_fp4", "quant_algo": quant_algo} @@ -842,6 +847,10 @@ class ModelConfig: return "fp8" elif self.quantization == "modelopt_fp4": return "nvfp4" + elif self.quantization == "modelopt_mixed": + raise ValueError( + "modelopt_mixed is only supported for pre-quantized checkpoints." + ) elif self.quantization == "modelopt": # Auto-detect from model config quant_cfg = self._parse_quant_hf_config() @@ -872,6 +881,7 @@ class ModelConfig: "modelopt", "modelopt_fp8", "modelopt_fp4", + "modelopt_mixed", ] modelopt_quantization_specified = ( self.quantization in _MODELOPT_QUANTIZATION_METHODS @@ -913,6 +923,7 @@ class ModelConfig: "marlin", "modelopt_fp8", "modelopt_fp4", + "modelopt_mixed", "gptq_marlin_24", "gptq_marlin", "awq_marlin", @@ -932,6 +943,7 @@ class ModelConfig: compatible_quantization_methods = { "modelopt_fp8": ["modelopt"], "modelopt_fp4": ["modelopt"], + "modelopt_mixed": ["modelopt"], "petit_nvfp4": ["modelopt"], "w8a8_int8": ["compressed-tensors", "compressed_tensors"], "w8a8_fp8": ["compressed-tensors", "compressed_tensors"], diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index 9174c08ff..8a6b1b06e 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -31,6 +31,7 @@ from sglang.srt.layers.quantization.gptq import GPTQConfig, GPTQMarlinConfig from sglang.srt.layers.quantization.modelopt_quant import ( ModelOptFp4Config, ModelOptFp8Config, + ModelOptMixedPrecisionConfig, ) from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config @@ -57,6 +58,7 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = { "modelopt": ModelOptFp8Config, # Auto-detect, defaults to FP8 "modelopt_fp8": ModelOptFp8Config, "modelopt_fp4": ModelOptFp4Config, + "modelopt_mixed": ModelOptMixedPrecisionConfig, "w8a8_int8": W8A8Int8Config, "w8a8_fp8": W8A8Fp8Config, "awq": AWQConfig, diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index e47305e45..dc4fa71fc 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -591,6 +591,183 @@ class ModelOptFp8KVCacheMethod(BaseKVCacheMethod): super().__init__(quant_config) +class ModelOptMixedPrecisionConfig(ModelOptQuantConfig): + """Configuration for ModelOpt MIXED_PRECISION checkpoints.""" + + def __init__( + self, + kv_cache_quant_algo: Optional[str], + exclude_modules: Optional[List[str]], + packed_modules_mapping: Optional[Dict[str, List[str]]], + quantized_layers: Dict[str, Dict[str, Any]], + fp8_config: ModelOptFp8Config, + nvfp4_config: "ModelOptFp4Config", + ) -> None: + super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping) + self.quantized_layers = quantized_layers + self.fp8_config = fp8_config + self.nvfp4_config = nvfp4_config + + @classmethod + def override_quantization_method(cls, hf_quant_config, user_quant): + if hf_quant_config is None: + return None + if hf_quant_config.get("quant_method", "") == "modelopt_mixed": + return "modelopt_mixed" + return None + + @classmethod + def get_name(cls) -> str: + return "modelopt_mixed" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return ModelOptFp4Config.get_min_capability() + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "ModelOptMixedPrecisionConfig": + kv_cache_quant_algo = None + exclude_modules = None + quantized_layers = {} + + quant_algo = config.get("quant_algo") + if quant_algo is not None: + kv_cache_scheme = config.get("kv_cache_scheme") + if isinstance(kv_cache_scheme, dict): + if ( + kv_cache_scheme.get("type") == "float" + and kv_cache_scheme.get("num_bits") == 8 + ): + kv_cache_quant_algo = "FP8" + elif ( + kv_cache_scheme.get("type") == "float" + and kv_cache_scheme.get("num_bits") == 4 + ): + kv_cache_quant_algo = "NVFP4" + else: + kv_cache_quant_algo = "auto" + exclude_modules = config.get("ignore") + quantized_layers = config.get("quantized_layers", {}) + else: + quantization_section = cls.get_from_keys(config, ["quantization"]) + quant_algo = quantization_section.get("quant_algo") + kv_cache_quant_algo = quantization_section.get("kv_cache_quant_algo") + exclude_modules = quantization_section.get("exclude_modules") + quantized_layers = quantization_section.get("quantized_layers", {}) + + if quant_algo != "MIXED_PRECISION": + raise ValueError( + "ModelOptMixedPrecisionConfig only supports MIXED_PRECISION checkpoints." + ) + if not quantized_layers: + raise ValueError( + "MIXED_PRECISION quantization requires a non-empty quantized_layers map." + ) + + group_size = None + for layer_info in quantized_layers.values(): + if layer_info.get("quant_algo", "").upper() == "NVFP4": + group_size = layer_info.get("group_size", 16) + break + if group_size is None: + group_size = 16 + + packed_modules_mapping = config.get("packed_modules_mapping") + fp8_config = ModelOptFp8Config( + is_checkpoint_fp8_serialized=True, + kv_cache_quant_method=kv_cache_quant_algo, + exclude_modules=[], + packed_modules_mapping=packed_modules_mapping, + ) + nvfp4_config = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=True, + kv_cache_quant_algo=kv_cache_quant_algo, + exclude_modules=[], + packed_modules_mapping=packed_modules_mapping, + group_size=group_size, + ) + + return cls( + kv_cache_quant_algo=kv_cache_quant_algo, + exclude_modules=exclude_modules, + packed_modules_mapping=packed_modules_mapping, + quantized_layers=quantized_layers, + fp8_config=fp8_config, + nvfp4_config=nvfp4_config, + ) + + def apply_weight_name_mapper(self, hf_to_sglang_mapper: "WeightsMapper"): + super().apply_weight_name_mapper(hf_to_sglang_mapper) + if self.quantized_layers: + self.quantized_layers = hf_to_sglang_mapper.apply_dict( + self.quantized_layers + ) + + def _resolve_quant_algo(self, prefix: str) -> Optional[str]: + if prefix in self.quantized_layers: + return self.quantized_layers[prefix]["quant_algo"].upper() + + proj_name = prefix.rsplit(".", 1)[-1] + if self.packed_modules_mapping and proj_name in self.packed_modules_mapping: + algos = set() + base = prefix.rsplit(".", 1)[0] + for shard_name in self.packed_modules_mapping[proj_name]: + shard_prefix = f"{base}.{shard_name}" + if shard_prefix in self.quantized_layers: + algos.add(self.quantized_layers[shard_prefix]["quant_algo"].upper()) + if len(algos) == 1: + return algos.pop() + if len(algos) > 1: + raise ValueError( + f"Mixed quant_algo within fused layer {prefix}: {algos}. " + "All shards must use the same quantization." + ) + + prefix_dot = prefix + "." + for key, info in self.quantized_layers.items(): + if key.startswith(prefix_dot): + return info["quant_algo"].upper() + + return None + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + quant_algo = self._resolve_quant_algo(prefix) + + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix, self.exclude_modules, self.packed_modules_mapping + ) or self.is_layer_excluded(prefix): + return UnquantizedLinearMethod() + if quant_algo == "FP8": + return ModelOptFp8LinearMethod(self.fp8_config) + if quant_algo == "NVFP4": + return ModelOptFp4LinearMethod(self.nvfp4_config) + return UnquantizedLinearMethod() + + if self.kv_cache_quant_algo and isinstance(layer, RadixAttention): + return ModelOptFp8KVCacheMethod(self.fp8_config) + + if isinstance(layer, FusedMoE): + if self.is_layer_excluded(prefix): + return None + if quant_algo == "FP8": + return ModelOptFp8MoEMethod(self.fp8_config) + if quant_algo == "NVFP4": + return ModelOptNvFp4FusedMoEMethod(self.nvfp4_config) + return None + + return None + + class ModelOptFp8MoEMethod(FusedMoEMethodBase): """MoE method for ModelOpt FP8. Supports loading FP8 checkpoints with static weight scale and activation scale. diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 3edcf14bd..576d02e3a 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -2728,7 +2728,8 @@ def get_model_loader( if model_config and ( (hasattr(model_config, "modelopt_quant") and model_config.modelopt_quant) - or model_config.quantization in ["modelopt_fp8", "modelopt_fp4", "modelopt"] + or model_config.quantization + in ["modelopt_fp8", "modelopt_fp4", "modelopt_mixed", "modelopt"] ): logger.info("Using ModelOptModelLoader due to ModelOpt quantization config.") return ModelOptModelLoader(load_config) @@ -2737,7 +2738,8 @@ def get_model_loader( if ( model_config and hasattr(model_config, "quantization") - and model_config.quantization in ["modelopt_fp8", "modelopt_fp4"] + and model_config.quantization + in ["modelopt_fp8", "modelopt_fp4", "modelopt_mixed"] ): if model_config._is_already_quantized(): logger.info( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2be1c045a..b5b11e562 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -105,6 +105,7 @@ QUANTIZATION_CHOICES = [ "modelopt", "modelopt_fp8", "modelopt_fp4", + "modelopt_mixed", "petit_nvfp4", "w8a8_int8", "w8a8_fp8", @@ -1546,7 +1547,8 @@ class ServerArgs: self.moe_a2a_backend == "none" and self.moe_runner_backend == "auto" and ( - self.quantization in ["fp8", "modelopt_fp8", "modelopt_fp4"] + self.quantization + in ["fp8", "modelopt_fp8", "modelopt_fp4", "modelopt_mixed"] or is_kimi_k2_k25_thinking_int4 ) ): @@ -1819,15 +1821,19 @@ class ServerArgs: "modelopt", "modelopt_fp8", "modelopt_fp4", + "modelopt_mixed", ]: assert model_config.hf_config.mlp_hidden_act == "relu2" if model_config.quantization == "modelopt": - self.quantization = ( - "modelopt_fp4" - if model_config.hf_config.quantization_config["quant_algo"] - == "NVFP4" - else "modelopt_fp8" - ) + quant_algo = model_config.hf_config.quantization_config[ + "quant_algo" + ] + if quant_algo == "MIXED_PRECISION": + self.quantization = "modelopt_mixed" + else: + self.quantization = ( + "modelopt_fp4" if quant_algo == "NVFP4" else "modelopt_fp8" + ) else: self.quantization = model_config.quantization self.moe_runner_backend = "flashinfer_cutlass" @@ -2479,8 +2485,9 @@ class ServerArgs: assert self.quantization in [ "modelopt_fp4", "modelopt_fp8", + "modelopt_mixed", None, - ], f"Invalid quantization '{self.quantization}'. \nFlashInfer Cutlass MOE supports only: 'modelopt_fp4', 'modelopt_fp8', or bfloat16 (None)." + ], f"Invalid quantization '{self.quantization}'. \nFlashInfer Cutlass MOE supports only: 'modelopt_fp4', 'modelopt_fp8', 'modelopt_mixed', or bfloat16 (None)." assert self.ep_size in [ 1, self.tp_size, @@ -2492,9 +2499,10 @@ class ServerArgs: "fp8", "mxfp8", "modelopt_fp8", + "modelopt_mixed", "compressed-tensors", None, - ], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'fp8', 'modelopt_fp8', 'compressed-tensors', or bfloat16 (None)." + ], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'fp8', 'modelopt_fp8', 'modelopt_mixed', 'compressed-tensors', or bfloat16 (None)." self.disable_shared_experts_fusion = True logger.warning( "FlashInfer TRTLLM MoE is enabled. --disable-shared-experts-fusion is automatically set." diff --git a/test/registered/model_loading/test_modelopt_loader.py b/test/registered/model_loading/test_modelopt_loader.py index 97a6adf6d..2658a87a5 100644 --- a/test/registered/model_loading/test_modelopt_loader.py +++ b/test/registered/model_loading/test_modelopt_loader.py @@ -14,7 +14,11 @@ 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.modelopt_utils import QUANT_CFG_CHOICES +from sglang.srt.layers.quantization.modelopt_quant import ( + ModelOptMixedPrecisionConfig, +) from sglang.srt.model_loader.loader import ModelOptModelLoader +from sglang.srt.models.utils import WeightsMapper from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -620,5 +624,66 @@ class TestParseQuantHfConfig(CustomTestCase): self.assertNotIn("quant_algo", result) +class TestModelOptMixedPrecisionConfig(CustomTestCase): + def test_nemotron_mixed_precision_uses_modelopt_mixed(self): + model_config = ModelConfig.__new__(ModelConfig) + model_config.hf_config = MagicMock() + model_config.hf_config.model_type = "nemotron_h" + model_config.hf_config.architectures = ["NemotronHForCausalLM"] + + result = model_config._parse_modelopt_quant_config( + {"quantization": {"quant_algo": "MIXED_PRECISION"}} + ) + + self.assertEqual(result["quant_method"], "modelopt_mixed") + + def test_mixed_precision_override_does_not_hijack_w4afp8(self): + self.assertIsNone( + ModelOptMixedPrecisionConfig.override_quantization_method( + {"quant_method": "w4afp8", "quant_algo": "MIXED_PRECISION"}, + "w4afp8", + ) + ) + + def test_mixed_precision_uses_nvfp4_min_capability(self): + self.assertEqual(ModelOptMixedPrecisionConfig.get_min_capability(), 100) + + def test_mixed_precision_quant_layer_resolution_after_mapping(self): + quant_config = ModelOptMixedPrecisionConfig.from_config( + { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "backbone.layers.0.mixer.in_proj": {"quant_algo": "FP8"}, + "backbone.layers.1.mixer.experts.0.up_proj": { + "quant_algo": "NVFP4", + "group_size": 16, + }, + "backbone.layers.2.mixer.q_proj": {"quant_algo": "FP8"}, + "backbone.layers.2.mixer.k_proj": {"quant_algo": "FP8"}, + "backbone.layers.2.mixer.v_proj": {"quant_algo": "FP8"}, + }, + "packed_modules_mapping": { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + }, + } + ) + quant_config.apply_weight_name_mapper( + WeightsMapper(orig_to_new_prefix={"backbone.": "model."}) + ) + + self.assertEqual( + quant_config._resolve_quant_algo("model.layers.0.mixer.in_proj"), + "FP8", + ) + self.assertEqual( + quant_config._resolve_quant_algo("model.layers.1.mixer.experts"), + "NVFP4", + ) + self.assertEqual( + quant_config._resolve_quant_algo("model.layers.2.mixer.qkv_proj"), + "FP8", + ) + + if __name__ == "__main__": unittest.main()