quant: extract shared checkpoint quant metadata resolver (#35172)

python/sglang/srt/model_loader/checkpoint_quantization.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Pure-data helpers for quantization metadata in Hugging Face configs."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any, Literal, Mapping, TypeAlias
+
+__all__ = [
+    "CheckpointQuantSpec",
+    "QuantMetadataSource",
+    "resolve_checkpoint_quant_spec",
+]
+
+
+QuantMetadataSource: TypeAlias = Literal[
+    "quantization_config",
+    "text_config.quantization_config",
+    "compression_config",
+]
+
+
+@dataclass(slots=True)
+class CheckpointQuantSpec:
+    """Quantization metadata declared by a checkpoint.
+
+    ``declared_method`` preserves ``quant_method`` verbatim and is never inferred
+    from backend-specific fields. This intentionally contains no runtime
+    quantization classes, model construction, or layer hierarchy.
+    """
+
+    declared_method: str | None
+    config: dict[str, Any]
+    source: QuantMetadataSource
+
+
+def _get_field(config: object, name: str) -> Any:
+    if isinstance(config, Mapping):
+        return config.get(name)
+    return getattr(config, name, None)
+
+
+def _to_metadata_dict(value: object, source: QuantMetadataSource) -> dict[str, Any]:
+    if isinstance(value, Mapping):
+        return deepcopy(dict(value))
+
+    to_dict = getattr(value, "to_dict", None)
+    if callable(to_dict):
+        metadata = to_dict()
+        if isinstance(metadata, Mapping):
+            return deepcopy(dict(metadata))
+
+    raise TypeError(
+        f"{source} must be a mapping or expose to_dict(), "
+        f"got {type(value).__name__}"
+    )
+
+
+def _select_hf_quant_metadata(
+    hf_config: object,
+) -> tuple[QuantMetadataSource, object] | None:
+    value = _get_field(hf_config, "quantization_config")
+    if value is not None:
+        return "quantization_config", value
+
+    text_config = _get_field(hf_config, "text_config")
+    value = _get_field(text_config, "quantization_config")
+    if value is not None:
+        return "text_config.quantization_config", value
+
+    value = _get_field(hf_config, "compression_config")
+    if value is not None:
+        return "compression_config", value
+
+    return None
+
+
+def resolve_checkpoint_quant_spec(hf_config: object) -> CheckpointQuantSpec | None:
+    """Resolve checkpoint quantization metadata from an HF config.
+
+    The lookup order matches SRT's checkpoint loader: top-level
+    ``quantization_config``, the text sub-config used by some multimodal
+    checkpoints, then ``compression_config``. The returned metadata is deep-copied
+    so callers can attach runtime-only fields without mutating the HF config.
+    """
+
+    selected = _select_hf_quant_metadata(hf_config)
+    if selected is None:
+        return None
+
+    source, value = selected
+    config = _to_metadata_dict(value, source)
+    declared_method = config.get("quant_method")
+    return CheckpointQuantSpec(
+        declared_method=(declared_method if isinstance(declared_method, str) else None),
+        config=config,
+        source=source,
+    )

python/sglang/srt/model_loader/weight_utils.py

@@ -42,15 +42,16 @@ from tqdm.auto import tqdm
 
 from sglang.srt.configs.load_config import LoadConfig
 from sglang.srt.configs.model_config import REQUANTIZATION_METHODS, ModelConfig
-from sglang.srt.distributed import (
+from sglang.srt.distributed import get_world_group
-    get_world_group,
-)
 from sglang.srt.layers.quantization import QuantizationConfig, get_quantization_config
 from sglang.srt.layers.quantization.fp8 import Fp8Config
 from sglang.srt.layers.quantization.modelopt_quant import (
     ModelOptFp4Config,
     ModelOptFp8Config,
 )
+from sglang.srt.model_loader.checkpoint_quantization import (
+    resolve_checkpoint_quant_spec,
+)
 from sglang.srt.model_loader.ci_weight_validation import (
     ci_download_with_validation_and_retry,
     ci_validate_and_cleanup_local_snapshot,
@@ -271,18 +272,9 @@ def get_quant_config(
     if model_config.quantization == "gguf":
         return quant_cls.from_config({})
 
-    # Read the quantization config from the HF model config, if available.
+    checkpoint_quant_spec = resolve_checkpoint_quant_spec(model_config.hf_config)
-    hf_quant_config = getattr(model_config.hf_config, "quantization_config", None)
+    if checkpoint_quant_spec is not None:
-    # some vision model may keep quantization_config in their text_config
+        hf_quant_config = checkpoint_quant_spec.config
-    hf_text_config = getattr(model_config.hf_config, "text_config", None)
-    if hf_quant_config is None and hf_text_config is not None:
-        hf_quant_config = getattr(hf_text_config, "quantization_config", None)
-    if hf_quant_config is None:
-        # compressed-tensors uses a compressions_config
-        hf_quant_config = getattr(model_config.hf_config, "compression_config", None)
-    if hf_quant_config is not None:
-        if not isinstance(hf_quant_config, dict):
-            hf_quant_config = hf_quant_config.to_dict()
         # For modelopt_mixed, config.json's quantization_config may not
         # contain all runtime metadata. Fall through to the file-based
         # hf_quant_config.json path when the per-layer map or KV-cache

test/registered/unit/model_loader/test_modelopt_loader.py

@@ -5,7 +5,10 @@ This test module verifies the functionality of ModelOptModelLoader, which
 applies NVIDIA Model Optimizer quantization to models during loading.
 """
 
+import json
+import tempfile
 import unittest
+from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
@@ -693,6 +696,98 @@ class TestModelOptFp4LoaderSelection(CustomTestCase):
 
 
 class TestModelOptMixedPrecisionConfig(CustomTestCase):
+    def test_incomplete_inline_config_falls_back_to_hf_quant_config_file(self):
+        packed_modules_mapping = {
+            "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        }
+        file_quantized_layers = {
+            "model.layers.0.self_attn.q_proj": {"quant_algo": "FP8"}
+        }
+        file_config = {
+            "producer": {"name": "modelopt"},
+            "quantization": {
+                "quant_algo": "MIXED_PRECISION",
+                "kv_cache_quant_algo": "FP8",
+                "exclude_modules": [],
+                "quantized_layers": file_quantized_layers,
+            },
+        }
+        inline_configs = (
+            {
+                "quant_method": "modelopt_mixed",
+                "quant_algo": "MIXED_PRECISION",
+                "kv_cache_quant_algo": "NVFP4",
+            },
+            {
+                "quant_method": "modelopt_mixed",
+                "quant_algo": "MIXED_PRECISION",
+                "quantized_layers": {
+                    "inline.layer": {"quant_algo": "NVFP4", "group_size": 16}
+                },
+            },
+        )
+
+        with tempfile.TemporaryDirectory() as model_path:
+            Path(model_path, "hf_quant_config.json").write_text(
+                json.dumps(file_config), encoding="utf-8"
+            )
+            for inline_config in inline_configs:
+                with self.subTest(inline_config=inline_config):
+                    model_config = SimpleNamespace(
+                        quantization="modelopt_mixed",
+                        hf_config=SimpleNamespace(
+                            quantization_config=inline_config,
+                        ),
+                        model_path=model_path,
+                        revision=None,
+                        is_draft_model=False,
+                        is_draft_quantization_explicit=False,
+                    )
+
+                    config = get_quant_config(
+                        model_config, LoadConfig(), packed_modules_mapping
+                    )
+
+                    self.assertIsInstance(config, ModelOptMixedPrecisionConfig)
+                    self.assertEqual(config.quantized_layers, file_quantized_layers)
+                    self.assertEqual(config.kv_cache_quant_algo, "FP8")
+                    self.assertEqual(
+                        config.packed_modules_mapping, packed_modules_mapping
+                    )
+
+    @patch("sglang.srt.model_loader.weight_utils.snapshot_download")
+    def test_complete_inline_config_does_not_download_metadata(self, mock_download):
+        packed_modules_mapping = {
+            "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        }
+        inline_quantized_layers = {
+            "model.layers.0.self_attn.q_proj": {"quant_algo": "FP8"}
+        }
+        model_config = SimpleNamespace(
+            quantization="modelopt_mixed",
+            hf_config=SimpleNamespace(
+                quantization_config={
+                    "quant_method": "modelopt_mixed",
+                    "quant_algo": "MIXED_PRECISION",
+                    "kv_cache_scheme": {"type": "float", "num_bits": 8},
+                    "exclude_modules": [],
+                    "quantized_layers": inline_quantized_layers,
+                }
+            ),
+            model_path="remote/model",
+            revision=None,
+            is_draft_model=False,
+            is_draft_quantization_explicit=False,
+        )
+
+        config = get_quant_config(model_config, LoadConfig(), packed_modules_mapping)
+
+        self.assertIsInstance(config, ModelOptMixedPrecisionConfig)
+        self.assertEqual(config.quantized_layers, inline_quantized_layers)
+        self.assertEqual(config.kv_cache_quant_algo, "FP8")
+        self.assertEqual(config.packed_modules_mapping, packed_modules_mapping)
+        mock_download.assert_not_called()
+
     def test_minimax_mixed_precision_resolves_runtime_names_and_mxfp8(self):
         quant_config = ModelOptMixedPrecisionConfig.from_config(
             {

test/registered/unit/test_checkpoint_quantization.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+
+from sglang.srt.model_loader.checkpoint_quantization import (
+    CheckpointQuantSpec,
+    resolve_checkpoint_quant_spec,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class _ConfigObject:
+    def __init__(self, **values):
+        self.__dict__.update(values)
+
+
+class _QuantConfigObject:
+    def __init__(self, values):
+        self._values = values
+
+    def to_dict(self):
+        return self._values
+
+
+class TestResolveCheckpointQuantSpec(CustomTestCase):
+    def test_top_level_quantization_config(self):
+        config = {
+            "quantization_config": {
+                "quant_method": "fp8",
+                "activation_scheme": "dynamic",
+            }
+        }
+
+        spec = resolve_checkpoint_quant_spec(config)
+
+        self.assertEqual(
+            spec,
+            CheckpointQuantSpec(
+                declared_method="fp8",
+                config={"quant_method": "fp8", "activation_scheme": "dynamic"},
+                source="quantization_config",
+            ),
+        )
+
+    def test_text_config_fallback_supports_config_objects(self):
+        config = _ConfigObject(
+            text_config=_ConfigObject(
+                quantization_config={"quant_method": "gptq", "bits": 4}
+            ),
+            compression_config={"quant_method": "compressed-tensors"},
+        )
+
+        spec = resolve_checkpoint_quant_spec(config)
+
+        self.assertIsNotNone(spec)
+        self.assertEqual(spec.declared_method, "gptq")
+        self.assertEqual(spec.source, "text_config.quantization_config")
+
+    def test_compression_config_fallback(self):
+        config = _ConfigObject(
+            compression_config={"quant_method": "compressed-tensors"}
+        )
+
+        spec = resolve_checkpoint_quant_spec(config)
+
+        self.assertIsNotNone(spec)
+        self.assertEqual(spec.declared_method, "compressed-tensors")
+        self.assertEqual(spec.source, "compression_config")
+
+    def test_modelopt_quant_algo_does_not_infer_declared_method(self):
+        config = {
+            "quantization_config": {
+                "quant_algo": "FP8",
+                "exclude_modules": ["lm_head"],
+            }
+        }
+
+        spec = resolve_checkpoint_quant_spec(config)
+
+        self.assertIsNotNone(spec)
+        self.assertIsNone(spec.declared_method)
+        self.assertEqual(spec.config["quant_algo"], "FP8")
+
+    def test_quant_config_object_is_converted(self):
+        config = _ConfigObject(
+            quantization_config=_QuantConfigObject(
+                {"quant_method": "bitsandbytes", "load_in_4bit": True}
+            )
+        )
+
+        spec = resolve_checkpoint_quant_spec(config)
+
+        self.assertIsNotNone(spec)
+        self.assertEqual(spec.config["load_in_4bit"], True)
+
+    def test_lookup_priority_matches_srt_loader(self):
+        config = {
+            "quantization_config": {},
+            "text_config": {"quantization_config": {"quant_method": "gptq"}},
+            "compression_config": {"quant_method": "compressed-tensors"},
+        }
+
+        spec = resolve_checkpoint_quant_spec(config)
+
+        self.assertIsNotNone(spec)
+        self.assertEqual(spec.config, {})
+        self.assertEqual(spec.source, "quantization_config")
+
+    def test_metadata_is_deep_copied(self):
+        metadata = {"quant_method": "fp8", "modules_to_not_convert": ["lm_head"]}
+        spec = resolve_checkpoint_quant_spec({"quantization_config": metadata})
+
+        self.assertIsNotNone(spec)
+        spec.config["modules_to_not_convert"].append("embed_tokens")
+
+        self.assertEqual(metadata["modules_to_not_convert"], ["lm_head"])
+
+    def test_missing_metadata_returns_none(self):
+        self.assertIsNone(resolve_checkpoint_quant_spec({"model_type": "qwen3_vl"}))
+
+    def test_invalid_metadata_type_has_clear_error(self):
+        with self.assertRaisesRegex(TypeError, "quantization_config must be a mapping"):
+            resolve_checkpoint_quant_spec({"quantization_config": "fp8"})
+
+
+if __name__ == "__main__":
+    unittest.main()