Hugging Face's logo

Diffusers documentation

加速推理

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

加速推理

扩散模型在推理阶段速度较慢，因为生成是一个迭代过程，需要在若干“步”中逐步将噪声细化为图像或视频。为了加速这一过程，您可以尝试使用不同的 调度器、降低模型权重的精度以加快计算、采用更节省显存的注意力机制等。

将这些技术组合使用，可以实现比单独使用任何一项技术更快的推理速度。

本指南将介绍如何加速推理。

模型数据类型

模型权重的精度和数据类型会影响推理速度，因为更高的精度需要更多的内存来加载，并且需要更多的时间来进行计算。PyTorch 默认以 float32 或全精度加载模型权重，因此更改数据类型是一种快速提升推理速度的简单方法。

bfloat16
float16
TensorFloat-32

bfloat16 与 float16 类似，但对数值误差的鲁棒性更强。bfloat16 的硬件支持情况因设备而异，但大多数现代 GPU 均支持 bfloat16。

import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

缩放点积注意力

内存高效注意力机制针对推理速度 和 内存占用 进行了优化！

缩放点积注意力（SDPA）实现了多种注意力后端，包括 FlashAttention、xFormers 以及原生 C++ 实现。它会自动为您的硬件选择最优的后端。

如果您使用的是 PyTorch >= 2.0，SDPA 默认已启用，且无需对代码进行任何额外更改。不过，如果您希望自行选择，也可以尝试使用其他注意力后端。下面的示例使用 torch.nn.attention.sdpa_kernel 上下文管理器来启用高效注意力。

from torch.nn.attention import SDPBackend, sdpa_kernel
import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
  image = pipeline(prompt, num_inference_steps=30).images[0]

torch.compile

torch.compile 通过将 PyTorch 代码和运算编译为优化后的内核来加速推理。Diffusers 通常会对计算量较大的模型（如 UNet、transformer 或 VAE）进行编译。

启用以下编译器设置以获得最大速度（有关更多选项，请参阅完整列表）。

import torch
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

加载并编译 UNet 和 VAE。您可以从多种不同的模式中进行选择，其中 "max-autotune" 通过编译为 CUDA graph 来优化至最快速度。CUDA graph 通过单次 CPU 操作启动多个 GPU 操作，从而有效降低开销。

在 PyTorch 2.3.1 中，您可以控制 torch.compile 的缓存行为。这对于 "max-autotune" 等编译模式尤其有益，该模式会对多个编译标志执行网格搜索，以找到最优配置。有关详细信息，请参阅 torch.compile 中的编译时缓存 教程。

将内存布局更改为 channels_last 还可以优化内存和推理速度。

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.to(memory_format=torch.channels_last)
pipeline.vae.to(memory_format=torch.channels_last)
pipeline.unet = torch.compile(
    pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.decode = torch.compile(
    pipeline.vae.decode,
    mode="max-autotune",
    fullgraph=True
)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

首次编译速度较慢，但编译完成后速度会显著提升。请尽量仅将编译后的 pipeline 用于同一类型的推理操作。如果对不同的图像尺寸调用编译后的 pipeline，会重新触发编译，这既缓慢又低效。

动态形状编译

请务必始终使用 PyTorch 的 nightly 版本，以获得更好的支持。

torch.compile 会跟踪输入的形状和条件，如果这些条件发生变化，它会重新编译模型。例如，如果模型是在 1024x1024 分辨率的图像上编译的，而随后用于不同分辨率的图像，则会触发重新编译。

为避免重新编译，请添加 dynamic=True，以尝试生成更具动态性的内核，从而在条件变化时避免重新编译。

+ torch.fx.experimental._config.use_duck_shape = False
+ pipeline.unet = torch.compile(
    pipeline.unet, fullgraph=True, dynamic=True
)

指定 use_duck_shape=False 可指示编译器是否应使用相同的符号变量来表示大小相同的输入。有关更多详细信息，请参阅此 评论。

并非所有模型都能开箱即用地从动态编译中获益，可能需要对其进行修改。请参阅此 PR，该 PR 改进了 AuraFlowPipeline 的实现，使其能够受益于动态编译。

如果动态编译在 Diffusers 模型上未按预期工作，请随时提交 issue。

区域编译

区域编译通过仅编译模型中较小且频繁重复的块（通常是 Transformer 层）来降低冷启动延迟，并允许在后续每次出现时复用编译产物。 对于许多扩散架构而言，这能带来与全图编译相同的运行时加速效果，并将编译时间缩短 8 到 10 倍。

使用 compile_repeated_blocks() 方法（一个封装了 torch.compile 的辅助函数）对诸如 transformer 模型等任意组件进行编译，如下所示。

# pip install -U diffusers
import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"

# compile only the repeated transformer layers inside the UNet
pipeline.unet.compile_repeated_blocks(fullgraph=True)

要为新模型启用区域编译，请在模型类中添加一个 _repeated_blocks 属性，其中包含要编译的块的类名（以字符串形式）。

class MyUNet(ModelMixin):
    _repeated_blocks = ("Transformer2DModel",)  # ← compiled by default

有关更多区域编译示例，请参阅参考 PR。

Accelerate 中还提供了一个 compile_regions 方法，可自动选择模型中需要编译的候选块。其余的图会被单独编译。由于无需手动设置要编译的块或调整编译标志，因此该方法非常适合快速实验。

# pip install -U accelerate
import torch
from diffusers import StableDiffusionXLPipeline
from accelerate.utils import compile_regions

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True)

compile_repeated_blocks() 的设计是刻意保持显式的。在 _repeated_blocks 中列出需要重复的块，该辅助函数只会编译这些块。它通过一行代码即可提供可预测的行为，并便于推理缓存的复用。

图断点

在 torch.compile 中指定 fullgraph=True 非常重要，以确保底层模型中不存在图中断。这样，您就可以充分利用 torch.compile，而不会导致任何性能下降。对于 UNet 和 VAE，这会改变您访问返回变量的方式。

- latents = unet(
-   latents, timestep=timestep, encoder_hidden_states=prompt_embeds
-).sample

+ latents = unet(
+   latents, timestep=timestep, encoder_hidden_states=prompt_embeds, return_dict=False
+)[0]

GPU 同步

step() 函数会在去噪器每次做出预测后被调用，并且会对 sigmas 变量进行索引。当将其放置在 GPU 上时，由于 CPU 与 GPU 之间的通信同步，会引入延迟。当去噪器已经被编译后，这一现象会更加明显。

通常，sigmas 应保留在 CPU 上，以避免通信同步和延迟。

Refer to the torch.compile and Diffusers: A Hands-On Guide to Peak Performance blog post for maximizing performance with torch.compile for diffusion models.

基准测试

请参阅 diffusers/benchmarks 数据集，查看编译后流水线的推理延迟和内存使用数据。

diffusers-torchao 仓库还包含 Flux 和 CogVideoX 编译版本的基准测试结果。

内核

Kernels 是一个用于构建、分发和加载 Hub 上优化计算内核的库。它支持 attention 内核以及用于 RMSNorm、GEGLU、RoPE 和 AdaLN 等操作的自定义 CUDA 内核。

Diffusers Pipeline Integration 指南展示了如何将内核与 add cuda-kernels 技能进行集成。该技能可让 Claude 或 Codex 等智能体针对特定模型和您的硬件编写自定义内核。

安装 add cuda-kernels 技能，以教会智能体如何编写内核。Codex 和 Claude 的通用自定义内核 博客文章对此进行了更详细的介绍。

例如，一个自定义的 RMSNorm 内核（由 add cuda-kernels 技能生成）配合 torch.compile 使用，可在 H100 上将 LTX-Video 的生成速度提升 1.43 倍。

动态量化

动态量化通过降低精度以执行更快的数学运算，从而提升推理速度。这种量化方式在运行时根据数据动态确定激活值的缩放方式，而不是使用固定的缩放因子。因此，缩放因子能够更精确地与数据对齐。

以下示例使用 torchao 库对 UNet 和 VAE 应用 动态 int8 量化。

请参阅我们的 torchao 文档，以了解如何使用 Diffusers 的 torchao 集成。

配置编译器标签以实现最大速度。

import torch
from torchao import apply_dynamic_quant
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True
torch._inductor.config.force_fuse_int_mm_with_mul = True
torch._inductor.config.use_mixed_mm = True

使用 dynamic_quant_filter_fn 过滤掉 UNet 和 VAE 中那些无法从动态量化中获益的部分线性层。

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

apply_dynamic_quant(pipeline.unet, dynamic_quant_filter_fn)
apply_dynamic_quant(pipeline.vae, dynamic_quant_filter_fn)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

融合投影矩阵

fuse_qkv_projections 方法目前处于实验阶段，其支持范围主要限于 Stable Diffusion 流水线。有关如何为其他流水线启用该功能的更多信息，请参阅此 PR。

在注意力块中，输入被投影到三个子空间，分别由投影矩阵 Q、K 和 V 表示。这些投影通常分别计算，但你可以将它们水平合并为一个矩阵，从而在单步中完成投影。这会增加输入投影矩阵乘法的规模，同时提升量化的效果。

pipeline.fuse_qkv_projections()

资源

Update on GitHub

←基本性能 缓存→