通過低精度量化優化用于圖像編輯的 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 的全部潛力直接引入您的桌面。

?