Modal 如何让 FLUX.1-dev 推理提速三倍

“从 API 获取 AI”的时代正迅速走向终结。

高质量的开放权重模型与高性能开源软件相结合，意味着你可以轻松运行自己的 API 来生成 图像、音乐或文本，并享有自托管带来的全部控制权和定制能力。

但拥有运行自己的生成式推理的能力会引出一系列问题：什么时候这样做才有意义，该怎么做，而且更重要的是，如何做到与专有生成式 API 相同的性能和品质？

我们最近分享了关于运行自己的 LLM 推理的结果和建议。但我们也喜欢媒体生成模型，而优化方式看起来大不相同。LLM 推理的关键在于找到合适的高层框架并调好各种参数，而基于扩散的图像模型则需要更贴近底层硬件。

在这篇博文中，我们将介绍我们如何将 Black Forest Labs 广受欢迎的 FLUX.1-dev 模型作为 Modal 上的自动扩缩容服务来运行，通过将推理速度提升三倍并加快冷启动，使其在速度和价格上可与专有服务商竞争。你可以在这里找到代码。

tl;dr：优化编译器和硬件感知带来 1.5 倍提升，近似缓存带来 2 倍提升

我们确定，要与提供 FLUX.1-dev 图像服务的 API 竞争，我们需要在三秒内返回结果。

应用“标准”优化（运行 Torch 编译器、切换数据布局、融合 QKV 计算）让我们完成了目标的一半。

然后我们应用了一种有趣的近似激活缓存技术——First Block Caching，将延迟再次减半。

实现基线

在开始提升性能之前，你需要先干净地测量当前性能。

我们从标准的 Hugging Face diffusers 库开始，并以 16 位精度创建我们的 FluxPipeline。

对各种输入取平均后，我们发现生成一张 1024x1024 的图像大约需要 6.75 秒。

应用标准优化，获得 1.5 倍加速

我们首先应用了一系列“标准”优化——使用 Torch 的优化编译器；融合 Transformer 注意力中的 query、key 和 value 计算；以及使用“channels last”内存布局。这些几乎总是好主意。

使用 Torch 编译器优化计算图

PyTorch 程序究竟是什么？默认情况下，PyTorch 在 Python 中动态构建张量运算的计算图，并以 eager 模式运行。这个“虚拟计算图”在主机/CPU 上执行，并触发设备/GPU 上“真实计算图”的执行。如果你好奇其工作原理，我们建议生成一些 PyTorch trace 并加以研究。

程序的图表示非常适合进行程序变换。在 BC 时代（ChatGPT 之前），大多数人用 PyTorch 训练自己的神经网络，而关键的程序变换是反向运行程序，以弄清楚如何让它不那么错，也就是“通过反向传播误差来学习表示”。

如今，就像神经网络本身一样，PyTorch 更多地被用于推理，因此关键的程序变换已转变为编译。编译将计算图替换为等价但更快的计算图。如果你熟悉数据库查询编译器，可以将其类比为谓词下推之类的逻辑-逻辑优化变换。

与任何像样的现代编译器一样，Torch 编译器以一系列逐步降低到越来越具体的中间表示的形式运行。TorchDynamo 挂钩 CPython 帧解释器，追踪 Python 字节码，并将一段段 Tensor 操作切分出来，降低为 “FX”图。随后，像 TorchInductor 这样的后端编译器会接收这些图，并将它们降低为更进一步优化的表示，例如 Triton 内核。

我们分别编译模型的两个大型子组件（Transformer 和变分自编码器）。我们的配置设置见下方代码片段。其中许多设置来自 Hugging Face 上的这份优秀指南，并在调整其他参数之前做了一些轻度验证。

最值得注意的选择是使用 max-autotune，这会带来数十分钟的编译时间成本，但能确保最优的运行时性能。关于我们如何在不失去 Modal 透明自动扩缩能力的情况下将其缩短到几分钟，请参见最后一节。

通过融合 QKV 向 GPU 暴露更多并行性

FLUX 包含一个大型 Transformer 模型。Transformer 架构的标志性组件是注意力块，它在文本与图像之间以及通过内部电路传递信息。它们通常以该块输入矩阵 X 与其权重矩阵 W_q、W_k 和 W_v 之间的三次独立矩阵乘法来表示。

如果我们将这三个权重矩阵拼接起来，就可以将注意力计算作为一次大型矩阵乘法来执行：(QKV = X @ W_qkv)。这会向降低后的表示暴露操作中更多的并行性。特别是，X 在整个乘法过程中都是同一个矩阵，而不是三个变量引用——我们完全保证 Torch 编译器必须验证它们指向完全相同的数据。

使用 channels-last 内存布局改善数据局部性

我们做的最后一项标准优化是改善数据局部性的常见建议。

张量是带劲的多维数组。表示图像（或图像上的特征图）的张量有三个维度：通道（颜色）、高度（y 位置）和宽度（x 位置）。这个三维数组需要映射到线性计算机内存上。

默认情况下，PyTorch 在内存中按通道优先、然后高度、然后宽度的顺序排列图像张量（CHW 或“Channels First”）。因此，顺序访问会从单个通道/颜色中读取空间上相邻的值。

让我们看一个例子。这张图像

在内存中以 CHW 格式表示为

但神经网络中的许多操作，例如卷积，是跨通道全局的、在空间上局部的。这意味着我们通常希望访问特定一组位置上的所有通道，因此我们希望通道排在最后。

你可以使用 memory_format 参数将 PyTorch 模型转换为这种格式。

综合起来，我们获得了 1.5 倍的加速

这些优化带来的总体性能提升约为 1.5 倍，主要由 Torch 编译器驱动。

这个加速效果绝对值得称道，下面的动画就证明了这一点，它展示了图像在去噪步骤中的演变过程，并以两种方法执行这些步骤的相同速率进行渲染。

应用基于感觉的近似缓存，再获得 2 倍加速

应用上述“标准”优化非常直接，最终成为工程投入/性能曲线上一个颇具吸引力的点。但我们需要在性能上走得更远，因此必须更深入。

扩散模型迭代式地生成图像，一步一步地将噪声转化为艺术。这就是我们在这些动画中展示的过程。如果你仔细观察，会发现某些步骤中图像几乎没有变化。

事实证明，如果你愿意容忍结果中的一些细微变化，你完全可以跳过那些步骤！

这是神经网络与其他程序之间的一个重要区别。对于神经网络，你通常可以移除大块内容或跳过步骤，程序仍然能运行，并且做“几乎”相同的事情。更像模拟计算机而非数字计算机！

我们使用了 ParaAttention 仓库中的“首块缓存”技术和实现，该仓库本身基于 TEACache 论文的方法。基本思路是开始为某个时间步运行模型。如果在模型前向传播的中途（在“首块”之后），看起来不会有大的变化，就跳过该步骤。

“大”的定义是一个可调参数，值越高会导致模型行为变化越大，但执行速度越快。与其他做同样事情的技术（如量化）相比，这允许在性能提升和质量下降之间实现更平滑的权衡。

我们在阈值为 0.12 时获得了 2 倍加速，且图像看起来比默认值 0.08 时更好，所以我们坚持使用了它。

通过缓存和快照将冷启动延迟降低 30 倍

在我们优化推理的过程中，启动时间遭受了巨大的打击——从几秒增加到几十分钟。

启动时间对成本和速度也很重要。如果启动速度快，你可以只运行满足当前需求所需的副本数量，同时仍能达到延迟目标。

我们认为这一点非常关键，并且在 Modal 上花了很多时间优化它！你可以在我们的 GPU 利用率解析和 与 Suno 的案例研究中了解更多关于为什么我们认为这对生成式应用如此重要的内容。

主要罪魁祸首是 Torch 编译器，具体来说是 max-autotune，它在编译时对多个实现进行性能分析以找到最快的那个。

这是缓存的经典用例——计算密集型工作产生可序列化的产物。Torch Compile 既提供了对较小产物（如编译后的 Triton 内核）的分段缓存，也提供了一个存储整个缓存计算图的“超级缓存”。我们两者都用了，但超级缓存并没有带来大幅加速。它也没有坏处，而且这是一个我们预期会随时间改进的新功能，所以我们保留了它。你可以在这里找到详细信息。

我们还使用 Modal 的内存快照节省了几秒钟，这让我们能够将 import torch 和 from_pretrained 中的大量文件读取和代码执行变成单次文件读取（对于第一次之后的每次调用）。查看这篇博客文章以深入了解。

使用 Modal 大规模服务 AI 模型

这些优化共同将 FLUX.1-dev 的服务延迟降低到与专有服务 API 相当的水平。在 Modal 上，这意味着你也能在价格上追平甚至击败这些提供商。

我们没有过多讨论构建和提供生成式 API 时出现的其他问题——交互式开发、处理突发负载，以及训练/评估服务的下一次迭代。如果你对此感兴趣，可以了解一下 Modal 无服务器平台，从 Suno 到 Substack 再到 足球队 的客户都信赖它，以数千块 GPU 和数万颗 CPU 的规模运行生成式推理。