[Diffusion] Add Diffusers quantization config bridge (#3837)

Signed-off-by: CAICAIIs <3360776475@qq.com>

vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py

@@ -0,0 +1,266 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantization helpers for the Diffusers backend."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from diffusers import PipelineQuantizationConfig, TorchAoConfig
+
+logger = logging.getLogger(__name__)
+
+_DIFFUSERS_TRANSFORMER_COMPONENTS = ("transformer", "transformer_2")
+_QuantizationConfigSpec = tuple[Callable[[Any], None], str]
+
+
+def ensure_supported_diffusers_quantization(quant_config: Any) -> None:
+    """Validate that a vLLM-Omni quantization config has a Diffusers mapping."""
+
+    method = _get_quant_method_name(quant_config)
+    validator, _ = _get_quantization_config_spec(method)
+    validator(quant_config)
+
+
+def convert_diffusers_quantization_config(load_kwargs: dict[str, Any]) -> None:
+    """Build Diffusers-native quantization config objects from dict kwargs."""
+
+    quant_config = load_kwargs.get("quantization_config")
+    if isinstance(quant_config, dict):
+        load_kwargs["quantization_config"] = _build_diffusers_native_quantization_config(quant_config)
+
+
+def apply_diffusers_quantization_config(
+    od_config: Any,
+    load_kwargs: dict[str, Any],
+    component_names: dict[str, Any],
+) -> bool:
+    """Inject a courtesy-converted quantization_config into load kwargs.
+
+    ``diffusers_load_kwargs`` is the canonical Diffusers backend configuration
+    path, so an explicit ``quantization_config`` already present in ``load_kwargs``
+    is never replaced. Returns whether a config was injected by this helper.
+    """
+
+    quant_config = getattr(od_config, "quantization_config", None)
+    if quant_config is None:
+        return False
+
+    if "quantization_config" in load_kwargs:
+        logger.warning(
+            "Both vLLM-Omni quantization_config and diffusers_load_kwargs.quantization_config "
+            "were provided for the diffusers backend. Using the Diffusers-native "
+            "quantization_config from diffusers_load_kwargs."
+        )
+        return False
+
+    transformer_components = _get_diffusers_transformer_components(component_names)
+    load_kwargs["quantization_config"] = _build_diffusers_quantization_config(
+        quant_config,
+        transformer_components,
+    )
+    return True
+
+
+def _normalize_method_name(method: Any) -> str:
+    method = getattr(method, "value", method)
+    return str(method).lower().replace("-", "_")
+
+
+def _get_quant_method_name(quant_config: Any) -> str:
+    get_name = getattr(quant_config, "get_name", None)
+    if get_name is not None:
+        return _normalize_method_name(get_name() if callable(get_name) else get_name)
+
+    method = getattr(quant_config, "quant_method", None)
+    if method is None:
+        method = getattr(quant_config, "method", None)
+    if method is None:
+        raise NotImplementedError(
+            "Diffusers backend quantization conversion requires a quantization "
+            "config with get_name(), quant_method, or method."
+        )
+    return _normalize_method_name(method)
+
+
+def _ensure_no_ignored_layers(quant_config: Any) -> None:
+    ignored_layers = getattr(quant_config, "ignored_layers", None)
+    if not ignored_layers:
+        ignored_layers = getattr(quant_config, "modules_to_not_convert", None)
+    if ignored_layers:
+        raise NotImplementedError(
+            "Diffusers backend quantization conversion does not map vLLM "
+            "ignored_layers/modules_to_not_convert names to Diffusers module "
+            "names. Use diffusers_load_kwargs for a native Diffusers config."
+        )
+
+
+def _get_torchao_quant_type_cls(class_name: str) -> type[Any]:
+    try:
+        import torchao.quantization as torchao_quantization
+    except ImportError as exc:
+        raise ImportError(
+            "Diffusers backend quantization conversion for fp8/int8 requires "
+            "torchao. Install torchao or pass a Diffusers-native "
+            "quantization_config through diffusers_load_kwargs."
+        ) from exc
+
+    try:
+        return getattr(torchao_quantization, class_name)
+    except AttributeError as exc:
+        raise ImportError(f"torchao.quantization.{class_name} is required for this quantization mapping.") from exc
+
+
+def _build_diffusers_native_quantization_config(config: dict[str, Any]) -> Any:
+    config = config.copy()
+    quant_mapping = config.get("quant_mapping")
+    if isinstance(quant_mapping, dict):
+        config["quant_mapping"] = {
+            component_name: _build_diffusers_component_quant_config(component_config)
+            for component_name, component_config in quant_mapping.items()
+        }
+    return PipelineQuantizationConfig(**config)
+
+
+def _build_diffusers_component_quant_config(config: Any) -> Any:
+    if not isinstance(config, dict):
+        return config
+
+    quant_method = config.get("quant_method")
+    if quant_method is None:
+        raise ValueError("Diffusers quant_mapping entries must provide quant_method.")
+
+    quant_method = _normalize_method_name(quant_method)
+    config_cls = _get_diffusers_quantization_config_cls(quant_method)
+    from_dict = getattr(config_cls, "from_dict", None)
+    if from_dict is not None:
+        return from_dict(config)
+
+    init_kwargs = config.copy()
+    init_kwargs.pop("quant_method", None)
+    return config_cls(**init_kwargs)
+
+
+def _get_diffusers_quantization_config_cls(quant_method: str) -> type[Any]:
+    from diffusers.quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING as DIFFUSERS_QUANT_CONFIG_MAPPING
+
+    config_cls = DIFFUSERS_QUANT_CONFIG_MAPPING.get(quant_method)
+    if config_cls is not None:
+        return config_cls
+
+    try:
+        from transformers.quantizers.auto import (
+            AUTO_QUANTIZATION_CONFIG_MAPPING as TRANSFORMERS_QUANT_CONFIG_MAPPING,
+        )
+    except ImportError:
+        TRANSFORMERS_QUANT_CONFIG_MAPPING = {}
+
+    config_cls = TRANSFORMERS_QUANT_CONFIG_MAPPING.get(quant_method)
+    if config_cls is None:
+        raise ValueError(f"Diffusers quant_mapping entry uses unknown quant_method={quant_method!r}.")
+    return config_cls
+
+
+def _build_torchao_pipeline_quant_config(
+    torchao_quant_type_name: str,
+    transformer_components: list[str],
+) -> Any:
+    """Build a Diffusers config after the caller has validated the vLLM config."""
+
+    quant_type_cls = _get_torchao_quant_type_cls(torchao_quant_type_name)
+    return PipelineQuantizationConfig(
+        quant_mapping={
+            component_name: TorchAoConfig(
+                quant_type=quant_type_cls(),
+            )
+            for component_name in transformer_components
+        }
+    )
+
+
+def _get_diffusers_transformer_components(component_names: dict[str, Any]) -> list[str]:
+    transformer_components = [
+        component_name for component_name in component_names if component_name in _DIFFUSERS_TRANSFORMER_COMPONENTS
+    ]
+    if transformer_components:
+        return transformer_components
+
+    available_components = ", ".join(sorted(component_names.keys())) or "<none>"
+    raise NotImplementedError(
+        "Diffusers backend quantization conversion currently only supports "
+        "pipelines with transformer or transformer_2 components. "
+        f"Found pipeline components: {available_components}. Use "
+        "diffusers_load_kwargs for a native Diffusers quantization config."
+    )
+
+
+def _validate_fp8_quant_config(quant_config: Any) -> None:
+    if getattr(quant_config, "is_checkpoint_fp8_serialized", False):
+        raise NotImplementedError(
+            "Diffusers backend fp8 conversion only supports online/dynamic "
+            "TorchAO quantization; serialized vLLM fp8 checkpoints are not mapped."
+        )
+
+    activation_scheme = getattr(quant_config, "activation_scheme", "dynamic")
+    if activation_scheme != "dynamic":
+        raise NotImplementedError(
+            f"Diffusers backend fp8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}."
+        )
+
+    weight_block_size = getattr(quant_config, "weight_block_size", None)
+    if weight_block_size is not None:
+        raise NotImplementedError(
+            "Diffusers backend fp8 conversion does not map vLLM weight_block_size "
+            "to TorchAO. Use diffusers_load_kwargs for a native Diffusers config."
+        )
+
+    _ensure_no_ignored_layers(quant_config)
+
+
+def _validate_int8_quant_config(quant_config: Any) -> None:
+    if getattr(quant_config, "is_checkpoint_int8_serialized", False):
+        raise NotImplementedError(
+            "Diffusers backend int8 conversion only supports online/dynamic "
+            "TorchAO quantization; serialized vLLM int8 checkpoints are not mapped."
+        )
+
+    activation_scheme = getattr(quant_config, "activation_scheme", "dynamic")
+    if activation_scheme != "dynamic":
+        raise NotImplementedError(
+            f"Diffusers backend int8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}."
+        )
+
+    _ensure_no_ignored_layers(quant_config)
+
+
+_QUANTIZATION_CONFIG_SPECS: dict[str, _QuantizationConfigSpec] = {
+    "fp8": (_validate_fp8_quant_config, "Float8DynamicActivationFloat8WeightConfig"),
+    "int8": (_validate_int8_quant_config, "Int8DynamicActivationInt8WeightConfig"),
+}
+
+
+def _get_quantization_config_spec(method: str) -> _QuantizationConfigSpec:
+    spec = _QUANTIZATION_CONFIG_SPECS.get(method)
+    if spec is None:
+        raise NotImplementedError(
+            f"Diffusers backend quantization conversion does not support {method!r}. "
+            "Use diffusers_load_kwargs for a native Diffusers quantization config, "
+            "or use a native vLLM-Omni pipeline for this quantization method."
+        )
+    return spec
+
+
+def _build_diffusers_quantization_config(
+    quant_config: Any,
+    transformer_components: list[str],
+) -> Any:
+    """Build a Diffusers PipelineQuantizationConfig from a supported config."""
+
+    method = _get_quant_method_name(quant_config)
+    _, torchao_quant_type_name = _get_quantization_config_spec(method)
+    return _build_torchao_pipeline_quant_config(
+        torchao_quant_type_name,
+        transformer_components,
+    )