Diffusion, Fast：纯 PyTorch 原生加速文生图扩散模型

想象这样一个场景：你在本地跑 Stable Diffusion XL 生成一张图，盯着进度条等了整整 7 秒——而隔壁同事同样一张图，只用了 2.4 秒。同样的模型、同样的显卡，为什么差距这么大？答案就藏在 HuggingFace 开源的 diffusion-fast 项目里。

从 7 秒到 2.4 秒：一场精准的性能优化战

这个项目起源于 PyTorch 官方博客系列「Accelerating Generative AI」的第三篇文章。PyTorch 团队选择了一个硬核目标：在不做任何模型蒸馏、量化压缩或换模型的前提下，仅靠 PyTorch 2.x 的原生能力，把 Stable Diffusion XL 的推理延迟从 7.36 秒压缩到 2.43 秒，整整提升 3 倍。

这不是靠运气。PyTorch 团队在 A100 GPU（80GB，400W）上逐层解锁了五项优化技术，每一步都有清晰的性能数据支撑，最终形成了一套可复制的加速流水线：

优化步骤	技术	延迟	加速比
基准	FP32 + 默认 Attention	7.36s	1.0x
第1步	BFloat16 精度	4.63s	1.6x
第2步	SDPA 注意力	3.31s	2.2x
第3步	torch.compile 编译	2.54s	2.9x
第4步	QKV 融合投影	2.52s	2.9x
第5步	动态 int8 量化	2.43s	3.0x

这套方法并非 SDXL 专属。团队同样将同样的优化应用于 SSD-1B（加速 2.13 倍）、Stable Diffusion v1.5（2.10 倍）、PixArt-XL-2-1024-MS（2.22 倍），验证了技术的通用性。

第一刀：BFloat16 精度 — 零改动的加速起点

精度转换是最直接的优化手段。现代 GPU 有专用的 Tensor Core 单元，在 bf16/fp16 这类降低精度格式下计算速度远超 fp32。BFloat16（Brain Float 16）相比标准 fp16 拥有更宽的指数位，数值范围更大，能有效避免扩散模型推理过程中因精度不足产生的数值溢出问题。

在 diffusers 库中，切换到 bf16 只需要在加载 Pipeline 时指定 torch_dtype=torch.bfloat16 即可：

from diffusers import StableDiffusionXLPipeline
import torch

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.bfloat16
).to("cuda")

仅此一行，延迟从 7.36 秒降至 4.63 秒，提升 37%。团队后续验证发现，torchao 新版本在 fp16 下也不会出现数值问题，如果对精度要求更高，可以改用 fp16 获得更快的速度。

第二刀：SDPA 注意力 — 让 GPU 做它最擅长的事

Scaled Dot Product Attention（SDPA）是 PyTorch 2.0 引入的 Flash Attention 的 PyTorch 原生实现。相比传统的手动实现，SDPA 能自动利用 GPU 的专用 Tensor Core，通过算子融合减少显存读写，在长序列场景下优势尤为明显。

在 diffusers 中，SDPA 是默认启用的——只要不手动调用 set_default_attn_processor() 覆盖掉它，SDPA 就会自动生效。这意味着在 bf16 基础上叠加 SDPA，几乎不需要修改代码，延迟从 4.63 秒进一步降至 3.31 秒。

第三刀：torch.compile — 编译出 GPU 最优执行计划

这是最关键的一步。torch.compile 是 PyTorch 2.0 的明星功能，它将 Python 代码编译成高度优化的 CUDA 内核，通过算子融合（operator fusion）和 CUDA Graphs 减少 GPU kernel 启动开销。

项目中同时编译了 UNet（去噪网络）和 VAE（解码器），并使用了 max-autotune 模式——让编译器搜索最优的执行配置：

# 转换为 channels_last 内存布局（融合卷积优化）
pipe.unet.to(memory_format=torch.channels_last)
pipe.vae.to(memory_format=torch.channels_last)

# 编译 UNet 和 VAE
pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True)
pipe.vae.decode = torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True)

关键参数是 fullgraph=True——强制编译全程不出现图断裂（graph break），确保整个计算图能被完整优化。若出现图断裂，部分子图会回退到解释执行，严重拖累性能。

为了确保 fullgraph 模式能通过，项目还手动做了两件事：把 conv2d 1x1 替换为等效的矩阵乘法（swap_conv2d_1x1_to_linear），以及在 max-autotune 模式下启用额外编译器参数（如 coordinate_descent_tuning）来微调 CUDA kernels。

第四刀：QKV 融合投影 — 减少三次矩阵乘法

标准扩散模型的 Attention 机制中，Q（Query）、K（Key）、V（Value）三个投影是分开计算的。fuse_qkv_projections() 将这三个矩阵拼接为一个大矩阵，一次矩阵乘法完成全部投影：

pipe.fuse_qkv_projections()

优势有两层：一是减少矩阵乘法次数（3 次到 1 次），二是投影后的矩阵更大，有利于后续 int8 量化的压缩效果。实测这一步带来了约 3.5% 的延迟改进。

第五刀：动态 int8 量化 — torchao 的精准刀法

这是最后一层优化，使用 PyTorch 官方性能优化库 torchao 实现动态量化。项目支持三种量化策略：

• int8 动态量化：apply_dynamic_quant，根据运行时数据范围动态缩放权重
• int8 权重仅量化：change_linear_weights_to_int8_woqtensors
• int4 权重仅量化：change_linear_weights_to_int4_woqtensors

关键在于 dynamic_quant_filter_fn——并非所有 Linear 层都适合量化。代码中排除了维度组合不利的层（如 (640, 1280)），确保量化层在精度和速度上都能受益：

def dynamic_quant_filter_fn(mod, *args):
    return (
        isinstance(mod, torch.nn.Linear)
        and mod.in_features > 16
        and (mod.in_features, mod.out_features) not in EXCLUDED_DIMS
    )

编译器也需相应调整以支持 int8 混合矩阵乘法：

torch._inductor.config.force_fuse_int_mm_with_mul = True
torch._inductor.config.use_mixed_mm = True

这一层优化将延迟从 2.52 秒压到 2.43 秒，约 3.5% 的提升。配合前四步，总加速比达到 3 倍。

代码结构：Benchmark 驱动的工程设计

diffusion-fast 不是推理服务，而是一套性能基准测试框架。代码结构清晰分为三层：

核心工具层（utils/）：

• pipeline_utils.py：封装 Pipeline 加载逻辑，统一处理 SDXL / SD v1.5 / PixArt-Alpha 等不同模型的加载差异
• benchmarking_utils.py：提供基准测试框架——使用 torch.utils.benchmark.Timer 精确测量执行时间，自动收集 GPU 显存占用，输出 CSV 格式结果

执行脚本层（根目录）：

• run_benchmark.py：主基准测试入口，加载 Pipeline → 预热 → 测量 → 输出 CSV + 效果预览图
• run_benchmark_pixart.py：PixArt 专用变体
• run_profile.py：性能剖析脚本
• prepare_results.py：合并多组实验结果，生成汇总 CSV 和可视化图表

实验脚本层（experiment-scripts/）：

• run_sd.sh：SDXL 实验命令模板
• run_sd_cpu.sh：CPU 推理实验（4 代 Intel Xeon 上 SDXL 也能获得 4 倍加速）
• run_pixart.sh：PixArt 实验命令

整个项目由 Dockerfile 提供可复现环境，基于 nvidia/cuda:12.1.0-runtime-ubuntu20.04，预装 PyTorch nightly 版本、diffusers 0.25.0 和 torchao，配合 Python 3.8 venv 确保依赖隔离。

适用场景与局限

适合的场景：

• 高频调用文生图 API 的开发者，需要对每次推理做成本控制
• 研究扩散模型推理优化的学者，需要一套可量化的性能基准
• 企业部署 AI 图像生成服务，需要榨干 GPU 硬件投资

局限性：

• 这不是推理服务：没有 Web UI、没有 API 接口，是一个 Python 基准测试工具包，需要自行集成到业务系统
• 需要兼容 GPU：优化效果主要在 NVIDIA Ampere+ 架构（A100/4090/3090）上明显，CPU 场景加速有限
• 依赖 nightly PyTorch：Dockerfile 中使用了带 CUDA 12.1 的 PyTorch nightly 版本（torch==2.3.0.dev20231218+cu121），稳定性不如稳定版，可能随版本更新 API 变化
• torch.compile 冷启动：编译本身有数十秒开销，适合批量推理而非单次生成

技术全景图

维度	内容
仓库	huggingface/diffusion-fast
主要语言	Python
核心框架	PyTorch 2.x, HuggingFace Diffusers
优化库	torchao
支持模型	SDXL, SD v1.5, SSD-1B, PixArt-Alpha
License	Apache-2.0
硬件需求	NVIDIA GPU (8GB+ VRAM)，CUDA 12.1+
基准测试	torch.utils.benchmark
量化方案	int8/int4 动态/权重量化
编译方案	torch.compile + CUDA Graphs
CPU 支持	支持（Intel Xeon 4代，4x 加速）

行业意义

diffusion-fast 的意义不仅在于 3 倍加速本身，而在于它展示了纯 PyTorch 原生能力的边界。通过五层递进优化，项目证明了在不借助 TensorRT、DeepSpeed-Inference 或模型蒸馏的前提下，PyTorch 2.x 已经足够强大。

这对 AI 开发者有直接启发：推理优化不必从换模型开始，bf16 + SDPA + compile 的组合拳往往能在现有基础上直接带来 2-3 倍提升。而 torchao 的引入则预示了 PyTorch 官方对训练后量化（Post-Training Quantization）的系统化支持方向。

图1：HuggingFace 官方 Logo — diffusion-fast 由 HuggingFace 官方团队维护