[chore] harden checkpoint quantization metadata parsing (#36922)

python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py

@@ -4,6 +4,7 @@ from unittest import mock
 
 import torch
 from torch import nn
+from transformers import PretrainedConfig
 
 from sglang.multimodal_gen.configs.models.encoders.clip import CLIPVisionConfig
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
@@ -135,7 +136,7 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase):
 
 class TestImageEncoderNativeLoading(unittest.TestCase):
     def test_bnb4_uses_shared_transformers_path_and_image_precision(self):
-        component_config = SimpleNamespace(
+        component_config = PretrainedConfig(
             is_encoder_decoder=False,
             architectures=["CLIPVisionModelWithProjection"],
             quantization_config={
@@ -193,7 +194,7 @@ class TestImageEncoderNativeLoading(unittest.TestCase):
         )
 
     def test_explicit_offload_is_rejected_before_transformers_load(self):
-        component_config = SimpleNamespace(
+        component_config = PretrainedConfig(
             is_encoder_decoder=False,
             architectures=["ThirdPartyVisionModel"],
             quantization_config={"quant_method": "fp8"},
@@ -236,7 +237,7 @@ class TestImageEncoderNativeLoading(unittest.TestCase):
             def to(self, *args, **kwargs):
                 raise AssertionError("quantized component must not be moved again")
 
-        component_config = SimpleNamespace(
+        component_config = PretrainedConfig(
             is_encoder_decoder=False,
             architectures=["ThirdPartyVisionModel"],
             quantization_config={"quant_method": "fp8"},

python/sglang/srt/model_loader/checkpoint_quantization.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
-"""Pure-data helpers for quantization metadata in Hugging Face configs."""
+"""Read quantization metadata declared by Hugging Face configurations."""
 
 from __future__ import annotations
 
@@ -8,6 +8,8 @@ from copy import deepcopy
 from dataclasses import dataclass
 from typing import Any, Literal, Mapping, TypeAlias
 
+from transformers import PretrainedConfig
+
 __all__ = [
     "CheckpointQuantSpec",
     "QuantMetadataSource",
@@ -20,11 +22,12 @@ QuantMetadataSource: TypeAlias = Literal[
     "text_config.quantization_config",
     "compression_config",
 ]
+ConfigMapping: TypeAlias = Mapping[str, Any] | PretrainedConfig
 
 
 @dataclass(slots=True)
 class CheckpointQuantSpec:
-    """Quantization metadata declared by a checkpoint.
+    """Quantization metadata declared by a checkpoint configuration.
 
     ``declared_method`` preserves ``quant_method`` verbatim and is never inferred
     from backend-specific fields. This intentionally contains no runtime
@@ -36,54 +39,48 @@ class CheckpointQuantSpec:
     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]:
+def _as_config_mapping(value: ConfigMapping, source: str) -> Mapping[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))
-
+        return value
+    if isinstance(value, PretrainedConfig):
+        return value.to_dict()
     raise TypeError(
-        f"{source} must be a mapping or expose to_dict(), "
+        f"{source} must be a mapping or transformers.PretrainedConfig, "
         f"got {type(value).__name__}"
     )
 
 
 def _select_hf_quant_metadata(
-    hf_config: object,
+    hf_config: ConfigMapping,
 ) -> tuple[QuantMetadataSource, object] | None:
-    value = _get_field(hf_config, "quantization_config")
+    config = _as_config_mapping(hf_config, "HF config")
+    value = config.get("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
+    text_config = config.get("text_config")
+    if text_config is not None:
+        text_config_mapping = _as_config_mapping(text_config, "text_config")
+        value = text_config_mapping.get("quantization_config")
+        if value is not None:
+            return "text_config.quantization_config", value
 
-    value = _get_field(hf_config, "compression_config")
+    value = config.get("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.
+def resolve_checkpoint_quant_spec(
+    hf_config: ConfigMapping,
+) -> CheckpointQuantSpec | None:
+    """Resolve quantization metadata from an HF configuration.
 
-    The lookup order matches SRT's checkpoint loader: top-level
+    The lookup order matches both serving runtimes: 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.
+    checkpoints, then ``compression_config``. Returned metadata is deep-copied
+    so callers can attach runtime-only fields without mutating the source config.
     """
 
     selected = _select_hf_quant_metadata(hf_config)
@@ -91,10 +88,11 @@ def resolve_checkpoint_quant_spec(hf_config: object) -> CheckpointQuantSpec | No
         return None
 
     source, value = selected
-    config = _to_metadata_dict(value, source)
-    declared_method = config.get("quant_method")
+    config = _as_config_mapping(value, source)
+    copied_config = deepcopy(dict(config))
+    declared_method = copied_config.get("quant_method")
     return CheckpointQuantSpec(
         declared_method=(declared_method if isinstance(declared_method, str) else None),
-        config=config,
+        config=copied_config,
         source=source,
     )