在我们之前的文章 diffusion-fast 中，我们展示了如何使用原生 PyTorch 代码将 Stable Diffusion XL (SDXL) 流水线的速度提升至原来的 3 倍。当时，SDXL 是开源领域最先进 (SoTA) 的图像生成流水线。不出所料，自那以后发生了很大变化，可以肯定地说，Flux 现已成为该领域能力最强的开放权重模型之一。
在这篇文章中，我们很高兴展示如何（主要）利用纯 PyTorch 代码并在 H100 等强悍 GPU 上，使 Flux.1-Schnell 和 Flux.1-Dev 的速度提升约 2.5 倍。
如果您迫不及待想要开始使用代码，可以在这里找到代码仓库。
优化概述
Diffusers 库中提供的流水线尽可能保持对 torch.compile 的友好支持。这意味着：
- 尽可能不出现图中断 (graph breaks)
- 尽可能避免重复编译
- 最少化或消除 CPU<->GPU 同步，以降低 Inductor 缓存查找开销
因此，它已经为我们提供了一个合理的起点。对于本项目，我们采用了与 diffusion-fast 项目相同的基本原则，并将其应用于 FluxPipeline。以下是我们所应用优化的概览（详情请见仓库）：
- 使用“fullgraph=True”和“max-autotune”模式进行
torch.compile，确保使用 CUDAgraphs - 合并 q,k,v 投影以进行注意力计算。这在量化过程中特别有用，因为它增加了维度，从而提高了计算密度
- 解码器输出使用
torch.channels_last内存格式 - Flash Attention v3 (FA3)，并将输入转换为（未缩放的）
torch.float8_e4m3fn - 动态 float8 激活量化以及通过 torchao
float8_dynamic_activation_float8_weight对线性层权重进行量化 - 用于在该模型上调整 Inductor 性能的一些标志：
- conv_1x1_as_mm = True
- epilogue_fusion = False
- coordinate_descent_tuning = True
- coordinate_descent_check_all_directions = True
- torch.export + 提前编译 Inductor (AOTI) + CUDAGraphs
除了以下两点外，大多数优化都不言自明：
- Inductor 标志。感兴趣的读者可以查看这篇博文以获取更多详细信息。
- 通过提前编译 (AoT compilation)，我们的目标是消除框架开销，并获得可以通过
torch.export导出的编译二进制文件。通过 CUDAGraphs，我们旨在优化内核启动。更多详情可在这篇文章中找到。
与大语言模型 (LLM) 不同，扩散模型主要受计算约束，因此 gpt-fast 中的优化并不能完全照搬到这里。下图展示了在 H100 700W GPU 上，每项优化（从左到右逐步叠加）对 Flux.1-Schnell 的影响：
对于 H100 上的 Flux.1-Dev，情况如下：
以下是应用不同优化后获得的 Flux.1-Dev 图像的视觉对比：
需要注意的是，只有 FP8 量化本质上是有损的，因此对于大多数此类优化，图像质量应保持一致。不过，在这种情况下，我们看到 FP8 带来的差异非常微小。
关于 CUDA 同步的说明
在调查过程中，我们发现去噪循环的第一步中，调度器中的这一步导致了一个 CPU<->GPU 同步点。我们可以通过在去噪循环开始时添加 self.scheduler.set_begin_index(0) 来消除它（PR）。
在使用 torch.compile 时，这个问题的影响更大，因为 CPU 必须等待同步完成才能进行 Dynamo 缓存查找，然后再在 GPU 上启动指令，而这个缓存查找过程稍显缓慢。因此，结论是：始终建议对您的流水线实现进行性能分析，并尽可能消除这些同步，以从编译中获益。
结论
本文介绍了一种使用原生 PyTorch 代码针对 Hopper 架构优化 Flux 的方案。该方案旨在平衡简洁性和性能。其他类型的优化也是可能的（例如使用融合 MLP 内核和融合自适应 LayerNorm 内核），但为了简洁起见，我们没有涉及这些内容。
另一个关键点是，Hopper 架构的 GPU 通常非常昂贵。因此，为了在消费级 GPU 上提供合理的“速度-内存”权衡，Diffusers 库中还有其他（通常兼容 torch.compile）的选项可供选择。我们邀请您在此处和此处查看它们。
我们欢迎您在其他模型上尝试这些技术并分享结果。优化愉快！