[diffusion] chore: reuse srt AutoRound for quantized DiTs (#36068)

python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py

@@ -2,6 +2,9 @@
 
 from typing import Literal, get_args
 
+from sglang.multimodal_gen.runtime.layers.quantization.auto_round import (
+    AutoRoundConfig,
+)
 from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
     BitsAndBytesConfig,
 )
@@ -27,6 +30,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.mxfp4_npu import (
 from sglang.multimodal_gen.runtime.layers.quantization.mxfp8 import MXFP8Config
 
 QuantizationMethods = Literal[
+    "auto-round",
     "fp8",
     "modelopt",
     "modelopt_fp8",
@@ -43,6 +47,7 @@ QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
 
 # The customized quantization methods which will be added to this dict.
 _CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {
+    "auto-round": AutoRoundConfig,
     "modelopt": ModelOptFp8DiffusionConfig,
     "modelopt_fp8": ModelOptFp8Config,
     "modelopt_fp4": ModelOptFp4Config,

python/sglang/multimodal_gen/runtime/layers/quantization/auto_round.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+from sglang.srt.layers.quantization.auto_round import AutoRoundConfig as SRTConfig
+
+
+class AutoRoundConfig(QuantizationConfig):
+    """Use SRT's serialized AutoRound kernels with diffusion linear layers."""
+
+    checkpoint_uses_native_qkv_layout = True
+
+    def __init__(self, srt_config: SRTConfig) -> None:
+        super().__init__()
+        self.srt_config = srt_config
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "auto-round"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return SRTConfig.get_supported_act_dtypes()
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return SRTConfig.get_min_capability()
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return SRTConfig.get_config_filenames()
+
+    @classmethod
+    def from_config(cls, config: dict) -> "AutoRoundConfig":
+        srt_config = SRTConfig.from_config(config)
+        if "gptq" not in srt_config.packing_format:
+            raise ValueError(
+                "SGLang diffusion currently supports AutoRound auto_gptq "
+                f"checkpoints, but got {srt_config.packing_format!r}."
+            )
+        return cls(srt_config)
+
+    def remap_checkpoint_prefixes(self, param_names_mapping: dict) -> None:
+        mapping = get_param_names_mapping(param_names_mapping)
+        remapped: dict[str, dict] = {}
+        for prefix, layer_config in (self.srt_config.extra_config or {}).items():
+            target, _, _ = mapping(f"{prefix}.weight")
+            target = target.removesuffix(".weight")
+            previous = remapped.setdefault(target, layer_config)
+            if previous != layer_config:
+                raise ValueError(
+                    f"AutoRound fused module {target!r} has inconsistent shard configs."
+                )
+
+        self.srt_config.extra_config = remapped
+        self.srt_config.block_name_to_quantize = None
+        self.srt_config.packed_modules_mapping = self.packed_modules_mapping
+
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
+        if not isinstance(layer, LinearBase):
+            return None
+
+        weight_bits, _, _ = self.srt_config.get_layer_config(layer, prefix)
+        if not self.srt_config.check_quantized(weight_bits):
+            return UnquantizedLinearMethod()
+
+        return self.srt_config.apply_gptq_quant_layer(
+            layer,
+            prefix,
+            self.srt_config.backend,
+            additional_linear_types=(LinearBase,),
+        )

python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py

@@ -45,3 +45,7 @@ class QuantizationConfig(SRTQuantizationConfig):
     ) -> bool:
         """Whether a row-parallel shard preserves this format's input layout."""
         return True
+
+    def remap_checkpoint_prefixes(self, param_names_mapping: dict) -> None:
+        """Translate checkpoint module names to the native model namespace."""
+        return

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -290,6 +290,15 @@ class TransformerLoader(ComponentLoader):
                 "Comfy quantized checkpoints do not support FSDP "
                 "inference; use TP and/or sequence parallelism instead"
             )
+        if (
+            use_fsdp
+            and quant_spec.quant_config is not None
+            and quant_spec.quant_config.get_name() == "auto-round"
+        ):
+            raise ValueError(
+                "AutoRound checkpoints do not support diffusion FSDP inference; "
+                "use TP and/or sequence parallelism instead"
+            )
 
         if quant_spec.gguf_file is not None:
             logger.info(

python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py

@@ -789,6 +789,9 @@ def resolve_transformer_quant_load_spec(
         packed = getattr(model_cls, "packed_modules_mapping", None)
         if packed and hasattr(quant_config, "packed_modules_mapping"):
             quant_config.packed_modules_mapping = packed
+        quant_config.remap_checkpoint_prefixes(
+            vars(model_cls).get("param_names_mapping", {})
+        )
 
     nunchaku_config = server_args.nunchaku_config
 
@@ -865,7 +868,7 @@ def _needs_device_weight_postprocess(
 ) -> bool:
     """Return whether post-load weight processing needs CUDA/NPU tensors."""
     quant_name = _get_quant_config_name(quant_config)
-    if quant_name in ("modelopt_fp8", "comfy_fp8", "mxfp8"):
+    if quant_name in ("modelopt_fp8", "comfy_fp8", "auto-round", "mxfp8"):
         return True
     if quant_name == "kitchen_int8":
         assert isinstance(quant_config, KitchenInt8Config)

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

@@ -50,6 +50,9 @@ from sglang.multimodal_gen.runtime.layers.linear import (
     ReplicatedLinear,
     UnquantizedLinearMethod,
 )
+from sglang.multimodal_gen.runtime.layers.quantization.auto_round import (
+    AutoRoundConfig,
+)
 from sglang.multimodal_gen.runtime.layers.quantization.comfy_fp8 import (
     ComfyFp8Config,
     ComfyFullPrecisionFp8LinearMethod,
@@ -95,6 +98,7 @@ from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
 )
 from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
 from sglang.multimodal_gen.runtime.models.dits.flux import FluxSingleTransformerBlock
+from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import MiniMaxH3DiTModel
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability
 from sglang.multimodal_gen.runtime.utils.quantization_utils import (
@@ -138,6 +142,41 @@ def _make_quant_config(name: str, **attrs):
 
 
 class TestTransformerQuantHelpers(unittest.TestCase):
+    def test_autoround_config_is_inferred_and_remapped_to_native_prefixes(self):
+        layer_config = {
+            "bits": 4,
+            "group_size": 128,
+            "sym": True,
+            "data_type": "int",
+            "act_bits": 16,
+        }
+        metadata = {
+            "quant_method": "auto-round",
+            "packing_format": "auto_round:auto_gptq",
+            **layer_config,
+            "block_name_to_quantize": "transformer_blocks",
+            "extra_config": {
+                "context_embedder": {**layer_config, "bits": 16},
+                **{
+                    f"transformer_blocks.0.attn.to_{shard}": layer_config
+                    for shard in ("q", "k", "v")
+                },
+            },
+        }
+
+        config = get_quant_config(
+            {"quantization_config": metadata}, "/unused/component/path"
+        )
+        self.assertIsInstance(config, AutoRoundConfig)
+        config.remap_checkpoint_prefixes(MiniMaxH3DiTModel.param_names_mapping)
+        self.assertEqual(
+            config.srt_config.get_layer_config(object(), "condition_proj")[0], 16
+        )
+        self.assertEqual(
+            config.srt_config.get_layer_config(object(), "blocks.0.attn.qkv_proj")[0],
+            4,
+        )
+
     def test_mps_layerwise_load_uses_residency_api(self):
         server_args = SimpleNamespace(
             should_configure_layerwise_offload_for_lazy_component=lambda name: (

python/sglang/srt/layers/quantization/auto_round.py

@@ -434,7 +434,13 @@ class AutoRoundConfig(QuantizationConfig):
                 return AWQLinearMethod(quant_args)
         return None
 
-    def apply_gptq_quant_layer(self, layer, prefix: str, backend: str = "auto"):
+    def apply_gptq_quant_layer(
+        self,
+        layer,
+        prefix: str,
+        backend: str = "auto",
+        additional_linear_types: tuple[type[torch.nn.Module], ...] = (),
+    ):
         from sglang.srt.layers.linear import LinearBase
         from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
         from sglang.srt.layers.quantization.gptq import (
@@ -449,9 +455,12 @@ class AutoRoundConfig(QuantizationConfig):
         from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
         from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
 
+        linear_types = (LinearBase, ParallelLMHead, *additional_linear_types)
+        is_linear = isinstance(layer, linear_types)
+
         weight_bits, group_size, sym = self.get_layer_config(layer, prefix)
         if not self.check_quantized(weight_bits):
-            if isinstance(layer, (LinearBase, ParallelLMHead)):
+            if is_linear:
                 return UnquantizedLinearMethod()
             else:
                 return None
@@ -475,7 +484,7 @@ class AutoRoundConfig(QuantizationConfig):
                 layer.scheme = quant_args.get_moe_scheme(layer)
                 return GPTQMoEMethod(quant_args)
 
-            if isinstance(layer, (LinearBase, ParallelLMHead)):
+            if is_linear:
                 layer.scheme = quant_args.get_linear_scheme(layer)
                 return GPTQLinearMethod(quant_args)
 
@@ -494,7 +503,7 @@ class AutoRoundConfig(QuantizationConfig):
                 layer.scheme = quant_args.get_moe_scheme(layer)
                 return GPTQMoEMethod(quant_args)
 
-            if isinstance(layer, (LinearBase, ParallelLMHead)):
+            if is_linear:
                 layer.scheme = quant_args.get_linear_scheme(layer)
                 return GPTQLinearMethod(quant_args)
 
@@ -551,7 +560,7 @@ class AutoRoundConfig(QuantizationConfig):
             }
             return MoeWNA16Config.from_config(config).get_quant_method(layer, prefix)
 
-        if isinstance(layer, (LinearBase, ParallelLMHead)):
+        if is_linear:
             if use_marlin:
                 return GPTQMarlinLinearMethod(quant_args_marlin)
             else: