[Feature]: FP8 Quantization Support for DiT (#1034)

Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>

vllm_omni/diffusion/quantization/__init__.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantization support for diffusion models.
+
+This module provides a unified interface for quantizing diffusion transformers
+using various methods (FP8, etc.). It wraps vLLM's quantization infrastructure
+while allowing diffusion-model-specific defaults and optimizations.
+
+Example usage:
+    from vllm_omni.diffusion.quantization import (
+        get_diffusion_quant_config,
+        get_vllm_quant_config_for_layers,
+    )
+
+    # Create FP8 config for diffusion model
+    diff_config = get_diffusion_quant_config("fp8")
+
+    # Get vLLM config to pass to linear layers
+    vllm_config = get_vllm_quant_config_for_layers(diff_config)
+
+    # Use in model initialization
+    linear_layer = QKVParallelLinear(..., quant_config=vllm_config)
+"""
+
+from typing import TYPE_CHECKING
+
+from vllm.logger import init_logger
+
+from .base import DiffusionQuantizationConfig
+from .fp8 import DiffusionFp8Config
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import (
+        QuantizationConfig,
+    )
+
+logger = init_logger(__name__)
+
+# Registry of supported quantization methods
+# To add a new method, create a new config class and register it here
+_QUANT_CONFIG_REGISTRY: dict[str, type[DiffusionQuantizationConfig]] = {
+    "fp8": DiffusionFp8Config,
+}
+
+SUPPORTED_QUANTIZATION_METHODS = list(_QUANT_CONFIG_REGISTRY.keys())
+
+
+def get_diffusion_quant_config(
+    quantization: str | None,
+    **kwargs,
+) -> DiffusionQuantizationConfig | None:
+    """Factory function to create quantization config for diffusion models.
+
+    Args:
+        quantization: Quantization method name ("fp8", etc.) or None to disable
+        **kwargs: Method-specific parameters passed to the config constructor
+
+    Returns:
+        DiffusionQuantizationConfig instance or None if quantization is disabled
+
+    Raises:
+        ValueError: If the quantization method is not supported
+
+    Example:
+        # Default FP8 with dynamic activation scaling
+        config = get_diffusion_quant_config("fp8")
+
+        # FP8 with custom parameters
+        config = get_diffusion_quant_config(
+            "fp8",
+            activation_scheme="static",
+            ignored_layers=["proj_out"],
+        )
+    """
+    if quantization is None or quantization.lower() == "none":
+        return None
+
+    quantization = quantization.lower()
+    if quantization not in _QUANT_CONFIG_REGISTRY:
+        raise ValueError(
+            f"Unknown quantization method: {quantization!r}. Supported methods: {SUPPORTED_QUANTIZATION_METHODS}"
+        )
+
+    config_cls = _QUANT_CONFIG_REGISTRY[quantization]
+    logger.info("Creating diffusion quantization config: %s", quantization)
+    return config_cls(**kwargs)
+
+
+def get_vllm_quant_config_for_layers(
+    diffusion_quant_config: DiffusionQuantizationConfig | None,
+) -> "QuantizationConfig | None":
+    """Get the vLLM QuantizationConfig to pass to linear layers.
+
+    This extracts the underlying vLLM config from a DiffusionQuantizationConfig,
+    which can then be passed to vLLM linear layers (QKVParallelLinear, etc.).
+
+    Args:
+        diffusion_quant_config: The diffusion quantization config, or None
+
+    Returns:
+        vLLM QuantizationConfig instance, or None if input is None
+    """
+    if diffusion_quant_config is None:
+        return None
+    return diffusion_quant_config.get_vllm_quant_config()
+
+
+__all__ = [
+    "DiffusionQuantizationConfig",
+    "DiffusionFp8Config",
+    "get_diffusion_quant_config",
+    "get_vllm_quant_config_for_layers",
+    "SUPPORTED_QUANTIZATION_METHODS",
+]

vllm_omni/diffusion/quantization/base.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Base class for diffusion model quantization configurations."""
+
+from abc import ABC
+from typing import TYPE_CHECKING, ClassVar
+
+import torch
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import (
+        QuantizationConfig,
+    )
+
+
+class DiffusionQuantizationConfig(ABC):
+    """Base class for diffusion model quantization configurations.
+
+    This provides a thin wrapper over vLLM's quantization configs,
+    allowing diffusion-model-specific defaults and future extensibility.
+
+    Subclasses should:
+        - Set quant_config_cls to the vLLM QuantizationConfig class
+        - Call super().__init__() after creating self._vllm_config
+        - Optionally override get_name() and get_min_capability() if needed
+    """
+
+    # Subclasses should set this to the vLLM QuantizationConfig class
+    quant_config_cls: ClassVar[type["QuantizationConfig"] | None] = None
+
+    # The underlying vLLM config instance
+    _vllm_config: "QuantizationConfig | None" = None
+
+    @classmethod
+    def get_name(cls) -> str:
+        """Return the quantization method name (e.g., 'fp8', 'int8').
+
+        By default, delegates to the underlying vLLM config class.
+        """
+        if cls.quant_config_cls is not None:
+            return cls.quant_config_cls.get_name()
+        raise NotImplementedError("Subclass must set quant_config_cls or override get_name()")
+
+    def get_vllm_quant_config(self) -> "QuantizationConfig | None":
+        """Return the underlying vLLM QuantizationConfig for linear layers."""
+        return self._vllm_config
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        """Return supported activation dtypes."""
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        """Minimum GPU compute capability required.
+
+        By default, delegates to the underlying vLLM config class.
+        """
+        if cls.quant_config_cls is not None:
+            return cls.quant_config_cls.get_min_capability()
+        return 80  # Ampere default

vllm_omni/diffusion/quantization/fp8.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FP8 quantization config for diffusion transformers."""
+
+from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+from .base import DiffusionQuantizationConfig
+
+
+class DiffusionFp8Config(DiffusionQuantizationConfig):
+    """FP8 quantization config optimized for diffusion transformers.
+
+    Uses dynamic activation scaling (no calibration dataset needed) and
+    online weight quantization from BF16/FP16 checkpoints.
+
+    Device Compatibility:
+        - Turing (SM 75+): Weight-only FP8 via Marlin kernel
+        - Ada/Hopper (SM 89+): Full W8A8 FP8 with native hardware support
+
+    The kernel selection is automatic based on GPU capability.
+
+    Args:
+        activation_scheme: Activation quantization scheme.
+            - "dynamic": Per-token dynamic scaling (default, no calibration)
+            - "static": Single per-tensor scale (requires calibration)
+        weight_block_size: Block size for block-wise weight quantization.
+            Format: [block_n, block_k]. If None, uses per-tensor scaling.
+        ignored_layers: List of layer name patterns to skip quantization.
+    """
+
+    # Tight coupling with vLLM's Fp8Config - delegates get_name() and get_min_capability()
+    quant_config_cls = Fp8Config
+
+    def __init__(
+        self,
+        activation_scheme: str = "dynamic",
+        weight_block_size: list[int] | None = None,
+        ignored_layers: list[str] | None = None,
+    ):
+        self.activation_scheme = activation_scheme
+        self.weight_block_size = weight_block_size
+        self.ignored_layers = ignored_layers or []
+
+        # Create underlying vLLM FP8 config
+        self._vllm_config = Fp8Config(
+            is_checkpoint_fp8_serialized=False,  # Online quantization from BF16
+            activation_scheme=activation_scheme,
+            weight_block_size=weight_block_size,
+            ignored_layers=ignored_layers,
+        )