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

[Quantization] Support config.json quantization_config format, fix exclude_modules matching, and fix KV cache scale loading for Nemotron (#18546)

Browse Source 
Signed-off-by: root <dafrimi@nvidia.com>
This commit is contained in:
danielafrimi
2026-02-21 10:14:29 +02:00
committed by GitHub
parent e239f8aa85
commit 33c33a7de9
3 changed files with 100 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -5,6 +5,7 @@ import logging
from enum import IntEnum
from typing import TYPE_CHECKING, Any, Dict, List, Optional
import regex as re
import torch
from torch.nn.parameter import Parameter
@@ -295,6 +296,11 @@ class ModelOptQuantConfig(QuantizationConfig):
elif self.kv_cache_quant_algo and isinstance(layer, RadixAttention):
return ModelOptFp8KVCacheMethod(self)
elif isinstance(layer, FusedMoE):
# Check if MoE layer should be excluded from quantization
# (e.g., MTP layers that have no quantization scales in checkpoint)
if self.is_layer_excluded(prefix):
# Falls back to default unquantized MoE
return None
return Moe(self)
return None
@@ -321,6 +327,54 @@ class ModelOptQuantConfig(QuantizationConfig):
# Preserve order, drop duplicates.
self.exclude_modules = list(dict.fromkeys(expanded))
def is_layer_excluded(self, prefix: str) -> bool:
"""Check if a layer should be excluded from quantization.
Handles:
- Exact matches (e.g., "lm_head" matching prefix "lm_head")
- Glob-style wildcards (e.g., "mtp*" matching "mtp_layers")
- Part-by-part matching (split prefix on "." and check each part)
- language_model. prefix stripping for vision-language models
- Fused module patterns (e.g., "q_a_proj" in "fused_qkv_a_proj_with_mqa")
"""
if not self.exclude_modules:
return False
# Build prefix variants: some models wrap layers under "language_model."
prefixes_to_check = [prefix]
if prefix.startswith("language_model."):
prefixes_to_check.append(prefix.removeprefix("language_model."))
# Fused module patterns: the exclude list may reference a sub-component
# (e.g., "q_a_proj") that is fused into a combined parameter name
# (e.g., "fused_qkv_a_proj_with_mqa"). We check if the last segment of
# the exclude pattern is a substring of the last segment of the prefix.
fused_patterns = {"q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"}
for pattern in self.exclude_modules:
# Convert glob-style wildcard to regex (e.g., "mtp*" -> "mtp.*")
regex_str = pattern.replace(".", r"\.").replace("*", r".*")
for pfx in prefixes_to_check:
if re.fullmatch(regex_str, pfx):
return True
# Part-by-part check: handles wildcards like "mtp*" matching
pfx_parts = pfx.split(".")
for part in pfx_parts:
if re.fullmatch(regex_str, part):
return True
# Check fused patterns: if the last segment of the exclude pattern
# is a known fused component, check if it appears in the prefix's
# last segment (handles fused_qkv_a_proj_with_mqa containing q_a_proj)
pattern_tail = pattern.rsplit(".", maxsplit=1)[-1]
if pattern_tail in fused_patterns:
for pfx in prefixes_to_check:
if pattern_tail in pfx.rsplit(".", maxsplit=1)[-1]:
return True
return False
class ModelOptFp8Config(ModelOptQuantConfig):
"""Configuration for ModelOpt FP8 quantization, including serialization and compatibility checks."""
@@ -376,19 +430,19 @@ class ModelOptFp8Config(ModelOptQuantConfig):
quant_method = config.get("quant_algo")
if quant_method is not None:
# Flat format (config.json quantization_config)
# For kv_cache, check if kv_cache_scheme exists and extract algo
# Derive kv_cache quant from kv_cache_scheme dict
kv_cache_scheme = config.get("kv_cache_scheme")
if (
kv_cache_scheme
and kv_cache_scheme.get("type") == "float"
and kv_cache_scheme.get("num_bits") == 8
):
kv_cache_quant_method = "FP8"
if isinstance(kv_cache_scheme, dict):
if (
kv_cache_scheme.get("type") == "float"
and kv_cache_scheme.get("num_bits") == 8
):
kv_cache_quant_method = "FP8"
# Map 'ignore' field to 'exclude_modules'
exclude_modules = config.get("ignore")
else:
# Fall back to nested format (hf_quant_config.json - legacy format)
# Fall back to nested format (hf_quant_config.json - will be deprecated)
try:
quantization_section = cls.get_from_keys(config, ["quantization"])
quant_method = quantization_section.get("quant_algo")
@@ -417,18 +471,6 @@ class ModelOptFp8Config(ModelOptQuantConfig):
packed_modules_mapping=config.get("packed_modules_mapping"),
)
def is_layer_excluded(self, prefix: str) -> bool:
if len(self.exclude_modules) == 0:
return False
return any(
module in prefix
or (
prefix.startswith("language_model.")
and module in prefix.removeprefix("language_model.")
)
for module in self.exclude_modules
)
def get_quant_method(
self, layer: torch.nn.Module, prefix: str
) -> Optional[QuantizeMethodBase]:
@@ -960,30 +1002,26 @@ class ModelOptFp4Config(ModelOptQuantConfig):
quant_method = config.get("quant_algo")
if quant_method is not None:
# Flat format (config.json quantization_config)
# Note: FP4 models in config.json format may not have all the detailed fields
# that are present in hf_quant_config.json, so we need to handle defaults
kv_cache_quant_algo = config.get("kv_cache_quant_algo")
if not kv_cache_quant_algo:
# For config.json format, derive from kv_cache_scheme if available
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"
else:
kv_cache_quant_algo = "auto"
elif isinstance(kv_cache_scheme, str):
scheme_name = kv_cache_scheme.strip().upper()
if scheme_name in ("FP8", "FLOAT8"):
kv_cache_quant_algo = "FP8"
elif scheme_name in ("FP4", "FLOAT4", "NVFP4"):
kv_cache_quant_algo = "NVFP4"
else:
kv_cache_quant_algo = "auto"
# Derive kv_cache_quant_algo from kv_cache_scheme dict
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"
else:
kv_cache_quant_algo = "auto"
elif isinstance(kv_cache_scheme, str):
scheme_name = kv_cache_scheme.strip().upper()
if scheme_name in ("FP8", "FLOAT8"):
kv_cache_quant_algo = "FP8"
elif scheme_name in ("FP4", "FLOAT4", "NVFP4"):
kv_cache_quant_algo = "NVFP4"
else:
kv_cache_quant_algo = "auto"
else:
kv_cache_quant_algo = "auto"
group_size = config.get("group_size")
# If group_size is not at top level, try to extract from config_groups
@@ -1038,27 +1076,6 @@ class ModelOptFp4Config(ModelOptQuantConfig):
config.get("packed_modules_mapping"),
)
def is_layer_excluded(self, prefix: str):
import regex as re
fused_patterns = ["q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"]
prefix_split = prefix.split(".")
for pattern in self.exclude_modules:
regex_str = pattern.replace(".", r"\.").replace("*", r".*")
pattern_split = pattern.split(".")
if re.fullmatch(regex_str, prefix):
return True
elif (
pattern_split[-1] in fused_patterns
and pattern_split[-1] in prefix_split[-1]
):
# Check if the last part of the excluded pattern is contained in the last part of the prefix
# This handles fused modules like fused_qkv_a_proj_with_mqa that contain q_a_proj and kv_a_proj_with_mqa
# e.g., model.layers.{i}.self_attn.{fused_weight_name}
assert len(prefix_split) == 5 and len(pattern_split) == 5
return True
return False
def get_quant_method(self, layer: torch.nn.Module, prefix: str):
return self._get_quant_method(
layer,

21
python/sglang/srt/model_loader/weight_utils.py
View File

@@ -1214,17 +1214,22 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:
return remapped_name
possible_scale_names = [".k_scale", ".v_scale"]
modelopt_scale_names = [".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale"]
# Patterns where modelopt stores scales under k_proj/v_proj
# but the model expects them under attn (RadixAttention)
modelopt_attn_prefixes = [".self_attn.", ".mixer."]
for scale_name in possible_scale_names:
if name.endswith(scale_name):
# Check and remap the name based on modelopt scale names
if any(
modelopt_scale_name in name
for modelopt_scale_name in modelopt_scale_names
):
# Check if this is a modelopt-style scale under k_proj/v_proj
matched_prefix = None
for attn_prefix in modelopt_attn_prefixes:
if f"{attn_prefix}{scale_name[1]}_proj{scale_name}" in name:
matched_prefix = attn_prefix
break
if matched_prefix is not None:
remapped_name = name.replace(
f".self_attn.{scale_name[1]}_proj{scale_name}",
f".self_attn.attn{scale_name}",
f"{matched_prefix}{scale_name[1]}_proj{scale_name}",
f"{matched_prefix}attn{scale_name}",
)
else:
remapped_name = name.replace(scale_name, f".attn{scale_name}")

7
python/sglang/srt/models/nemotron_h.py
View File

@@ -61,6 +61,7 @@ from sglang.srt.model_loader.weight_utils import (
replace_prefix,
replace_substrings,
)
from sglang.srt.models.utils import WeightsMapper
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import (
add_prefix,
@@ -640,6 +641,12 @@ class NemotronHForCausalLM(nn.Module):
"v_proj.v_scale": "attn.v_scale",
}
hf_to_sglang_mapper = WeightsMapper(
orig_to_new_prefix={
"backbone.": "model.",
}
)
def __init__(
self,
*,