Cache-DiT 加速指南

本指南介绍了如何在 vLLM-Omni 中使用 cache-dit 加速来加快扩散模型推理速度。

概述

Cache-dit 是一个通过智能缓存机制加速扩散 Transformer 模型的库。它支持多种加速技术，这些技术可以组合使用以获得最佳性能。

• DBCache：双块缓存，用于减少冗余计算。
• TaylorSeer：基于泰勒展开的预测，实现更快的推理。
• SCM：步进计算掩码，用于选择性步进计算。

快速入门

基本用法

只需设置 cache_backend="cache_dit" 即可启用 cache-dit 加速。Cache-dit 将使用其推荐的默认参数。

from vllm_omni.entrypoints.omni import Omni

# Simplest way: just enable cache-dit with default parameters
omni = Omni(
    model="Qwen/Qwen-Image",
    cache_backend="cache_dit",
)

images = omni.generate(
    "a beautiful landscape",
    num_inference_steps=50,
)

默认参数：当未提供 cache_config 时，cache-dit 会使用优化的默认值。有关所有参数及其默认值的完整列表，请参阅配置参考部分。

自定义配置

要自定义 cache-dit 设置，请提供一个 cache_config 字典，例如：

omni = Omni(
    model="Qwen/Qwen-Image",
    cache_backend="cache_dit",
    cache_config={
        "Fn_compute_blocks": 1,
        "Bn_compute_blocks": 0,
        "max_warmup_steps": 4,
        "residual_diff_threshold": 0.12,
    },
)

在线服务 (兼容 OpenAI)

通过在启动服务器时传递 --cache-backend cache_dit 来为在线服务启用 Cache-DiT。

# Use Cache-DiT default (recommended) parameters
vllm serve Qwen/Qwen-Image --omni --port 8091 --cache-backend cache_dit

要为在线服务自定义 Cache-DiT 设置，请通过 --cache-config 传递 JSON 字符串。

vllm serve Qwen/Qwen-Image --omni --port 8091 \
  --cache-backend cache_dit \
  --cache-config '{"Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4, "residual_diff_threshold": 0.12}'

加速方法

有关详细说明，请参阅 cache-dit 的用户指南。

1. DBCache (双块缓存)

当连续步进之间的残差差异较小时，DBCache 会智能地缓存中间 Transformer 块的输出，从而在不牺牲质量的情况下减少冗余计算。

关键参数:

• Fn_compute_blocks (int, 默认值: 1): 用于计算稳定特征差异的**前 n**个 Transformer 块的数量。值越高，缓存决策越准确，但计算量会增加。
• Bn_compute_blocks (int, 默认值: 0): 用于额外融合的**后 n**个 Transformer 块的数量。这些块充当近似隐藏状态的自动调整器。
• max_warmup_steps (int, 默认值: 4): 缓存被禁用的初始步数，以确保模型在缓存开始前学习到足够的特征。针对少量步进的蒸馏模型进行了优化。
• residual_diff_threshold (float, 默认值: 0.24): 残差差异的阈值。值越高，性能越快，但精度可能会降低。默认值使用相对较高的阈值以实现更积极的缓存。
• max_cached_steps (int, 默认值: -1): 最大缓存步数。设置为 -1 表示无限制缓存。
• max_continuous_cached_steps (int, 默认值: 3): 最大连续缓存步数。限制连续缓存以防止精度下降。

示例配置:

cache_config={
    "Fn_compute_blocks": 8,      # Use first 8 blocks for difference computation
    "Bn_compute_blocks": 0,       # No additional fusion blocks
    "max_warmup_steps": 8,        # Cache after 8 warmup steps
    "residual_diff_threshold": 0.12,  # Higher threshold for faster inference
    "max_cached_steps": -1,        # No limit on cached steps
}

性能提示:

• 对于大多数情况，默认的 Fn_compute_blocks=1 效果很好。对于更大的模型或需要更高精度的情况，请将其增加到 8-12。
• 增加 residual_diff_threshold (例如，0.12-0.15) 可实现更快的推理，但质量略有权衡；或者从默认值 0.24 降低以获得更高质量。
• 默认的 max_warmup_steps=4 针对少量步进的模型进行了优化。如果需要，请将其增加到 6-8。

2. TaylorSeer

TaylorSeer 使用泰勒展开来预测未来的隐藏状态，使模型能够跳过一些计算步骤，同时保持质量。

关键参数:

• enable_taylorseer (bool, 默认值: False): 启用 TaylorSeer 加速。
• taylorseer_order (int, 默认值: 1): 泰勒展开的阶数。阶数越高，精度越好，但计算量也越大。

示例配置:

cache_config={
    "enable_taylorseer": True,
    "taylorseer_order": 1,  # First-order Taylor expansion
}

性能提示:

• 在大多数情况下使用 taylorseer_order=1 (速度和质量的良好平衡)。
• 与 DBCache 结合使用以实现最大的加速效果。
• 更高的阶数 (2-3) 可能会提高质量，但会降低速度优势。

3. SCM (步进计算掩码)

SCM 允许您指定必须计算的步数以及可以使用缓存结果的步数，类似于 LeMiCa/EasyCache 风格的加速。

关键参数:

• scm_steps_mask_policy (str | None, 默认值: None): 预定义的掩码策略。选项：
• None: SCM 禁用 (默认)。
• "slow": 计算步数较多，质量较高 (28 步中计算 18 步)。
• "medium": 平衡 (28 步中计算 15 步)。
• "fast": 缓存步数较多，推理速度更快 (28 步中计算 11 步)。
• "ultra": 最大速度 (28 步中计算 8 步)。
• scm_steps_policy (str, 默认值: "dynamic"): 缓存步数的策略。
• "dynamic": 为掩码步数使用动态缓存 (推荐)。
• "static": 为掩码步数使用静态缓存。

示例配置:

cache_config={
    "scm_steps_mask_policy": "medium",  # Balanced speed/quality
    "scm_steps_policy": "dynamic",      # Use dynamic cache
}

性能提示:

• SCM 默认禁用 (scm_steps_mask_policy=None)。如果您需要额外的加速，请通过设置策略值来启用它。
• 从 "medium" 策略开始，并根据质量要求进行调整。
• 当质量可以略微妥协时，使用 "fast" 或 "ultra" 以获得最大的速度。
• "dynamic" 策略通常比 "static" 具有更好的质量。
• 当推理期间 num_inference_steps 发生变化时，SCM 掩码会自动重新生成。

配置参考

DiffusionCacheConfig 参数

参数	类型	默认值	描述
Fn_compute_blocks	int	1	用于计算差异的前 n 个块 (针对单 Transformer 模型进行了优化)。
Bn_compute_blocks	int	0	用于融合的后 n 个块。
max_warmup_steps	int	4	缓存开始之前的步数 (针对少量步进的蒸馏模型进行了优化)。
max_cached_steps	int	-1	最大缓存步数 (-1 = 无限制)。
max_continuous_cached_steps	int	3	最大连续缓存步数 (防止精度下降)。
residual_diff_threshold	float	0.24	残差差异阈值 (值越高，缓存越积极)。
num_inference_steps	int | None	None	SCM 掩码生成的初始推理步数 (可选，在推理期间自动刷新)。
enable_taylorseer	bool	False	启用 TaylorSeer 加速 (不适用于少量步进的蒸馏模型)。
taylorseer_order	int	1	泰勒展开的阶数。
scm_steps_mask_policy	str | None	None	SCM 掩码策略 (None, "slow", "medium", "fast", "ultra")。
scm_steps_policy	str	"dynamic"	SCM 计算策略 ("dynamic" 或 "static")。

示例：使用 CacheDiT 加速文本到图像生成

有关完整的实际示例，请参阅 examples/offline_inference/text_to_image/text_to_image.py。

# Enable cache-dit with hybrid acceleration
cd examples/offline_inference/text_to_image
python text_to_image.py \
    --model Qwen/Qwen-Image \
    --prompt "a cup of coffee on the table" \
    --cache_backend cache_dit \
    --num_inference_steps 50

该脚本使用 cache-dit 加速，并结合了 DBCache、SCM 和 TaylorSeer 的混合配置。

omni = Omni(
    model="Qwen/Qwen-Image",
    cache_backend="cache_dit",
    cache_config={
        # Scheme: Hybrid DBCache + SCM + TaylorSeer
        # DBCache
        "Fn_compute_blocks": 8,
        "Bn_compute_blocks": 0,
        "max_warmup_steps": 4,
        "residual_diff_threshold": 0.12,
        # TaylorSeer
        "enable_taylorseer": True,
        "taylorseer_order": 1,
        # SCM
        "scm_steps_mask_policy": "fast",  # Set to None to disable SCM
        "scm_steps_policy": "dynamic",
    },
)

您可以修改 cache_config 字典来自定义配置，以仅使用特定方法（例如，仅 DBCache，DBCache + SCM 等），具体取决于您的质量和速度要求。

要测试其他模型，您可以修改 --model 参数为目标模型标识符，例如 Tongyi-MAI/Z-Image-Turbo，并根据模型架构（例如 Transformer 块的数量）更新 cache_config。

额外资源

• Cache-DiT 用户指南
• Cache-DiT 基准测试
• DBCache 技术细节