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

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