[Feature]: FP8 Quantization Support for DiT (#1034)
Signed-off-by: lishunyang <lishunyang12@163.com> Signed-off-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
This commit is contained in:
1 parent
f0bca6f66b
commit
830d261f9d
17 files changed
+688
-13
No files matched your search
@@ -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,
|
||||
)
|
||||
Reference in new issue
Block a user