b690a10e3a
Signed-off-by: Lancer <maruixiang6688@gmail.com> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
81 lines
3.3 KiB
Python
81 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
|
|
|
|
@dataclass
|
|
class MagCacheConfig:
|
|
"""
|
|
Configuration for MagCache applied to transformer models.
|
|
|
|
MagCache (Magnitude-based Cache) is an adaptive caching technique that speeds up
|
|
diffusion model inference by reusing transformer block computations based on
|
|
magnitude ratios between consecutive timesteps.
|
|
|
|
Reference: https://github.com/Zehong-Ma/MagCache
|
|
|
|
Args:
|
|
threshold: Accumulated error threshold. Higher = more aggressive skipping (faster, lower quality).
|
|
Default: 0.24
|
|
max_skip_steps: Max consecutive skip steps (K).
|
|
Default: 5
|
|
retention_ratio: Fraction of initial steps where skipping is disabled (stability).
|
|
Default: 0.1
|
|
num_inference_steps: Total inference steps. Required for retention step calculation.
|
|
Default: 28
|
|
mag_ratios: Pre-computed magnitude ratios per step. Calibrate or use strategy defaults.
|
|
Default: None
|
|
mag_calibrate: If True, runs without skipping and logs norm_ratios for calibration.
|
|
Default: False
|
|
transformer_type: Transformer class name for logging.
|
|
Default: "FluxTransformer2DModel"
|
|
"""
|
|
|
|
threshold: float = 0.24
|
|
max_skip_steps: int = 5
|
|
retention_ratio: float = 0.1
|
|
num_inference_steps: int = 28
|
|
mag_ratios: torch.Tensor | list[float] | None = None
|
|
mag_calibrate: bool = False
|
|
transformer_type: str = "FluxTransformer2DModel"
|
|
|
|
def __post_init__(self) -> None:
|
|
"""Validate and set default coefficients."""
|
|
if self.threshold <= 0:
|
|
raise ValueError(f"threshold must be positive, got {self.threshold}")
|
|
|
|
if self.max_skip_steps <= 0:
|
|
raise ValueError(f"max_skip_steps must be positive, got {self.max_skip_steps}")
|
|
|
|
if not 0 < self.retention_ratio < 1:
|
|
raise ValueError(f"retention_ratio must be in (0, 1), got {self.retention_ratio}")
|
|
|
|
if self.num_inference_steps is None:
|
|
raise ValueError(
|
|
"num_inference_steps must be provided for MagCache. "
|
|
"This is required to determine retention steps and interpolate mag_ratios. "
|
|
"For Flux models, use num_inference_steps=28."
|
|
)
|
|
|
|
if self.num_inference_steps <= 0:
|
|
raise ValueError(f"num_inference_steps must be positive, got {self.num_inference_steps}")
|
|
|
|
if not self.mag_calibrate and self.mag_ratios is None:
|
|
raise ValueError(
|
|
"mag_ratios must be provided for MagCache inference because these ratios "
|
|
"are model-dependent. To get them for your model:\n"
|
|
"1. Initialize MagCacheConfig(mag_calibrate=True, ...)\n"
|
|
"2. Run inference on your model once.\n"
|
|
"3. Copy the printed ratios array and pass it to mag_ratios in the config.\n"
|
|
"For Flux models, you can import FLUX_MAG_RATIOS from vllm_omni.diffusion.cache.magcache.strategy."
|
|
)
|
|
|
|
if not self.mag_calibrate and self.mag_ratios is not None:
|
|
if not torch.is_tensor(self.mag_ratios):
|
|
self.mag_ratios = torch.tensor(self.mag_ratios)
|