通过低精度量化优化用于图像编辑的 FLUX.1 Kontext

FLUX.1 Kontext 是 Black Forest Labs 最近发布的模型，是对社区图像生成模型的一项令人着迷的补充。本文重点介绍的开放式权重 FLUX.1 Kontext [dev] 变体是一个针对图像到图像转换任务精心优化的模型。这一开创性工具因其增量图像编辑功能而脱颖而出，在指导生成过程方面实现了范式转变。

与依赖制作复杂提示并集成难以获取的蒙版或深度图和边缘图的传统方法不同，FLUX.1 Kontext [dev] 为用户提供了一种更直观、更灵活的方法。通过将增量编辑与用于提高扩散模型推理性能的前沿优化技术无缝集成，这一新模型能够提供截然不同的图形编辑用户体验。

NVIDIA 与 Black Forest Labs 合作，使用 NVIDIA TensorRT 和量化优化 NVIDIA RTX GPU 的 FLUX.1 Kontext [dev]，以更低的显存要求提供更快的推理速度。基于 NVIDIA TensorRT 为 NVIDIA Blackwell GeForce RTX 50 系列 GPU 解锁 FP4 图像生成所展示的基础工作，本文将介绍低精度量化如何彻底改变用户体验。如需全面的模型描述，请参阅 Black Forest Labs 技术报告。

多轮图像编辑

传统上，图像到图像的模型依赖于文本提示，并辅以二进制蒙版或深度/ 边缘图以进行优化控制。但是，这种范式很麻烦，需要对多个输入进行细致的设计，最终会阻碍这些模型的广泛应用。

FLUX.1 Kontext [dev] 引入了一种新方法，用户可以主要根据提供的提示有选择地编辑图像 (图 1) 。这种创新允许将复杂的编辑任务无缝划分为多阶段流程。值得注意的是，该模型在所有阶段都保留了原始图像的语义完整性。

在 FLUX.1 Kontext [dev] 环境中，推理速度至关重要，因为更快的模型使用户能够在短时间内逐步测试新的更改。最终，基于聊天的全新用户体验将成为图像编辑的终极目标。另一方面，降低内存消耗也很重要，因为图像生成流程由多个模型组成。随着模型变得更加强大，所需的内存也会增加，从而限制其采用率。通过以较低精度量化模型，可以在更多 NVIDIA RTX 硬件上执行此类流程，从而提高采用率或降低推理成本。

与之前的工作一样，优化工作针对配备第五代 Tensor Core 并原生支持 4 位浮点 (FP4) 运算的 NVIDIA Blackwell GPU。用于这项工作的工具包括用于执行量化的 NVIDIA TensorRT Model Optimizer 和用于生成最终引擎的推理运行时的 TensorRT。

工作流和量化

本节将重点介绍管道的计算方面，分析关键组件和优化策略以提高性能。

FLUX.1 [dev] 模型由几个关键模块组成：视觉 Transformer 主干、自动编码器、CLIP 和 T5 (图 2) 。值得注意的是，FLUX.1 Kontext [dev]与其他变体 (例如 FLUX.1-Deph-dev) 的不同之处在于，它将输入图像转换为令牌，然后将这些令牌连接到扩散过程中使用的输出图像令牌。这种设计选择具有重大意义：与其他 FLUX.1-dev 变体相比，其上下文窗口长度几乎翻了一番，从而增加了内存和计算需求。早期的实验表明，Transformer 模块消耗了大约 96% 的总处理时间，使其成为理想的优化目标。

众所周知的 Transformer 架构主要包括通用矩阵乘法 (GEMM) 运算 (例如线性层中的运算) 和扩展的点积注意力 (SDPA) 机制。虽然大量研究侧重于 GEMM 在各种低精度设置中的量化技术，但对利用低精度格式进行注意力运算的关注较少。SageAttention 系列论文是一个值得注意的例外。

图 3 突出显示了 SDPA 和 GEMM 内核的计算影响，约占 Transformer 模块总计算预算的 75%。请注意注意力 (SPDA) 和线性层 (GEMM) 如何主导 Transformer 模块的总体计算预算。

为平衡准确性和速度，我们采用了以下量化策略进行 SDPA 优化：

• 使用量化配置中指定的每张量 FP8 或每块 FP4 量化，通过低精度矩阵乘法将查询、键和值标记投影到公共向量空间。
• 从第一个矩阵乘法获得的输出以更高的精度累积，并且需要额外的量化运算将第一个批量矩阵乘法 (BMM1) 的输入转换为 FP8。换言之，投影 QKT 在 FP8 中执行。
• Softmax 运算目前仅支持 FP32，因此在此阶段不执行量化。
• 第二批乘法 (BMM2) 也使用每张量 FP8 量化设置进行计算。

对于 FP8 和 FP4 变体，SPDA 的拟议量化模式相似，仅在用于查询、键和值预测的精度方面存在差异。

图 4 显示了此模式的图形表示。使用逐块 FP4 量化模式对输入进行量化。BMM 使用每张量量化在 FP8 中计算。Softmax 始终在 FP32 中执行。最后，请注意，图中报告的 Quant 层实际上是指 FakeQuantize 运算。

在众多量化方法中，工作重点是 FP8 中的后训练量化 (PTQ) 、FP4 中的 PTQ 和 FP4 中的 SVD-Quant。整个量化过程使用 TensorRT Model Optimizer 库执行，而 TensorRT 用作融合各种量化内核的编译器。整体流程如下：

1. 定义 TensorRT Model Optimizer 使用的适当量化配置。
2. 使用 TensorRT Model Optimizer，根据配置将量化层注入到原始 FLUX.1 Kontext [dev] 模型中。
3. 加载由提示词/ 条件图像对组成的校准数据集。
4. 使用 TensorRT Model Optimizer 执行训练后校准。
5. 将量化模型导出为开放神经网络交换 (ONNX) 模型格式。
6. 使用 TensorRT 加载 ONNX 模型并针对特定目标硬件执行内核融合，从而生成最终的 TensorRT 引擎。

模型输出和性能

表 1 报告了通过 (量化) 视觉转换器执行单个扩散步骤所需的毫秒数；因此值越小，效率越高。从 BF16 到 FP8 的性能提升是巨大的，因为 8 位精度降低了内存带宽需求并增加了计算吞吐量。与 FP8 相比，FP4 的增益较小，其原因有两个：

1. 图 3 展示了 Attention 如何在 Transformer 模块中成为主要运算符。上一节介绍了 SPDA 如何主要在 FP8 中进行量化，以保持数值稳定性；因此，与 FP8 相比，限制了 FP4 的端到端加速。
2. 上下文窗口的采用时间是其他 Flux1.dev 模型的两倍，这进一步影响了 FP4 的优势。FP8：Attention 算子的二次计算成本。

图 5 显示了 FLUX.1 Kontext [dev] 模型在不同数值精度下的显存使用情况。当仅关注 Transformer 主干 (图 5，右) 时，从 BF16 过渡到 FP8 和 FP4 时，低精度量化可分别节省大约 2 倍和 3 倍的内存。

但是，由于模型的其他组件未进行量化，完整工作流的整体显存减少并不显著。节省内存的模型特别有价值，因为它们可以更轻松地部署在消费级 GPU 上，例如 5090 D。同样，服务提供商可以利用减少的显存占用来运行同一模型的多个实例，从而提高吞吐量和成本效益。

最后，图 6 和图 7 显示了低精度模型与全精度模型相比的质量。

图 6。FLUX.1 Kontext【dev】模型在各种精度格式下的图像编辑性能，以提示“删除此图像中的所有人员”

总结

FLUX.1 Kontext [dev] 模型的加速可以带来响应灵敏的用户体验，部分原因在于其增量编辑功能。这种协同作用将迷人图像的生成转化为交互式过程，这将激励创意社区采用这种创新的内容创作方法。

从技术角度来看，这篇博文介绍了有关量化 scale-dot-product-attention 运算符 ( Transformer 架构的关键组件) 的 TensorRT Model Optimizer 方法的宝贵见解。研究 AI 模型的研究人员可以在此基础上开发出能够更有效地利用低精度数据类型的新技术。提高 Attention 算子的推理时间效率，同时保持较高的数值准确性，可能会对更广泛的机器学习社区产生巨大的有益影响。

同样，模型开发者可以将 TensorRT Model Optimizer 和 TensorRT 用作加速推理和简化模型部署的可靠工具。

最新 NVIDIA 硬件与 TensorRT 和 TensorRT Model Optimizer 的结合带来了强大的性能组合。这种协同作用实现了尖端 AI 技术的普及，使更广泛的用户群能够利用其强大功能。通过在云功能和本地工作站之间架起桥梁，我们很高兴能将生成式 AI 的全部潜力直接引入您的桌面。