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

[Model] Support Nemotron 3 Super NVFP4 (#20407)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-03-14 15:56:26 +08:00
committed by GitHub
parent c95dc88f86
commit 75a7879fd4
6 changed files with 277 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/configs/model_config.py
View File

@@ -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"],

2
python/sglang/srt/layers/quantization/__init__.py
View File

@@ -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,

177
python/sglang/srt/layers/quantization/modelopt_quant.py
View File

@@ -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.

6
python/sglang/srt/model_loader/loader.py
View File

@@ -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(

26
python/sglang/srt/server_args.py
View File

@@ -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."