FlashAttention-3是FlashAttention系列的最新版本，由Colfax Research、Meta、英偉達等多個機構的研究團隊於2024年7月聯合提出。該算法針對英偉達H100 GPU的Hopper架構新特性進行了專門最佳化，旨在充分利用硬體能力以提升大模型訓練中注意力機制的計算效率。
相比前代FlashAttention-2，FlashAttention-3在FP16精度下速度提升1.5-2倍，計算吞吐量高達740 TFLOPS，達到了H100理論最大FLOPS的75%；在FP8精度下速度接近1.2 PFLOPS，同時誤差比標準注意力基線減少2.6倍。為實現這些改進，FlashAttention-3主要採用了通過warp專門化技術重疊計算與數據移動、交錯進行分塊矩陣乘法與softmax運算、以及利用硬體支持進行FP8低精度非相干處理等關鍵技術。
研發背景,發展歷程,技術特點,核心技術,核心團隊,重要事件,
研發背景
FlashAttention-2在H100 GPU上僅實現了35%的理論最大FLOP利用率，未能充分利用硬體功能。因此，FlashAttention-3的目標是在Hopper架構GPU（如H100、H800）上充分利用WGMMA和TMA的異步性加速注意力計算，以實現更高效的GPU利用率、在較低精度（如FP8）下保持更好性能並支持大模型處理更長的上下文序列，例如將H100 GPU的利用率從35%提升至75%。
發展歷程
2022年，一種快速、記憶體高效的注意力算法FlashAttention問世。2023年，研究團隊推出FlashAttention-2。2024年7月，來自Meta、英偉達、Together AI等機構的研究者聯合推出FlashAttention-3。
其研發背景在於FlashAttention-2在H100GPU上僅實現了約35%的理論最大FLOP利用率。因此FlashAttention-3旨在充分利用Hopper架構（如H100）的WGMMA、TMA和FP8等新硬體特性，以實現更高效的GPU利用率、在低精度下保持更好性能，並支持處理更長上下文的語言模型。
技術特點
FlashAttention-3採用了三項主要技術以加速Hopper GPU上的注意力計算，包括利用Tensor Core和TMA異步性的生產者-消費者模型與warp專門化、GEMM與softmax操作的異步重疊執行（包括組間桌球調度與組內兩階段流水線），以及FP8低精度下的分塊量化與非相干處理技術。
FlashAttention-3引入了一種“生產者-消費者”的異步編程模型。其中，“生產者”利用張量記憶體加速器（TMA）將數據從高頻寬記憶體異步載入至共享記憶體。“消費者”則從共享記憶體讀取數據並使用Tensor Core執行矩陣乘法計算。通過warp專門化技術，不同的warp被分配為生產者或消費者角色，使兩者能夠並行執行。此外，採用桌球調度策略，讓一個warp組執行矩陣乘法時，另一個warp組執行softmax，從而實現計算與數據移動的重疊。
在warp組內部，FlashAttention-3通過兩階段GEMM-softmax流水線打破操作間的數據依賴，實現計算重疊。第一階段，當前疊代的softmax操作與下一疊代的Q·K^T矩陣乘法重疊執行。第二階段，當前疊代的P·V矩陣乘法與下一疊代的softmax操作重疊執行。該方案通過引入額外的暫存器與共享記憶體緩衝區（如S_cur和S_next）實現跨疊代的數據傳遞與重用。
為適配FP8 Tensor Core，FlashAttention-3引入了記憶體布局轉換技術，動態轉置V矩陣的塊以滿足其布局要求。在精度控制方面，採用了分塊量化策略，為每個數據塊設定獨立的縮放因子以更好地捕捉局部數值分布。同時，通過非相干處理技術（如使用Hadamard變換旋轉輸入數據）破壞塊間相干性，減少量化誤差的傳播。
在性能方面，FlashAttention-3在FP16精度下的速度是FlashAttention-2的1.5-2.0倍，計算吞吐量高達740 TFLOPS，達到了H100理論最大FLOPS的75%利用率；在FP8精度下，其速度接近1.2 PFLOPS。與標準Attention相比，FlashAttention-3的速度快了3-16倍；對於中長序列（1k以上），其性能甚至超過了專門為H100最佳化的cuDNN。在精度控制上，FP16模式下由於中間結果保留在FP32中，其均方根誤差（RMSE）比標準Attention減少1.7倍；FP8模式下，得益於塊量化和非相干處理，其誤差比基線FP8注意力減少2.6倍。
當前版本的FlashAttention-3專注於Hopper架構（如H100/H800 GPU），暫不支持其他GPU型號，且目前著重於訓練過程的最佳化。未來的工作方向包括提升推理性能，並計畫將技術推廣到Hopper GPU以外的其他硬體架構。
核心技術
FlashAttention-3充分利用了Hopper GPU的新硬體特性，如WGMMA、TMA和FP8，旨在實現更高效的GPU利用率、在較低精度下獲得更好的性能，並能夠在大模型中使用更長的上下文。
通過異步計算與數據重疊最佳化性能，FlashAttention-3引入“生產者-消費者”編程模型，將注意力計算劃分為數據生產和消費兩個角色，並採用warp專門化技術用不同的warp分別匹配生產者和消費者以實現並行執行。其“桌球調度”策略讓一個warp組執行矩陣乘法時另一個warp組執行softmax，從而實現計算的重疊，這利用了Tensor Core和TMA的異步性來重疊整體計算和數據移動。
通過計算重排打破依賴，FlashAttention-3提出兩階段GEMM-softmax流水線方案，實現warp組內部GEMM和softmax的跨疊代重疊執行，以繞過操作間的順序依賴。這種方式重新安排了計算的執行順序以提高GPU利用率。
利用硬體加速的低精度計算，FlashAttention-3使用FP8張量核心使計算吞吐量翻倍，並採用分塊量化策略為每個數據塊單獨設定縮放因子以更好捕捉局部數值分布。通過非相干處理技術，如使用Hadamard變換，用隨機正交矩陣旋轉輸入數據以分散異常值，減少量化誤差傳播。此外，還進行了記憶體布局最佳化，動態轉置數據塊以適應FP8 Tensor Core對輸入數據布局的特定要求。
這些技術使得FlashAttention-3在FP16下計算吞吐量達740 TFLOPS/s，達到H100理論最大FLOPS利用率的75%，在FP8下接近1.2 PFLOPS/s，且誤差比標準注意力減少2.6倍。
核心團隊
FlashAttention-3 由 Colfax Research、Meta、英偉達（NVIDIA）、喬治亞理工學院（Georgia Tech）、普林斯頓大學（Princeton University）和 Together AI 的研究團隊於 2024 年 7 月聯合提出。英偉達的CUTLASS團隊和cuDNN團隊為研究提供了直接支持。Tri Dao是FlashAttention-1/2/3系列的核心作者之一，也是Mamba的共同一作。他於2023年6月在史丹福大學獲得計算機博士學位，師從Christopher Ré和Stefano Ermon，畢業後擔任Together AI的首席科學家，並於2024年6月入職普林斯頓大學擔任助理教授。此外，Meta、Together AI 和普林斯頓大學為該項研究提供了計算支持。
重要事件
2024年7月（具體為7月12日），FlashAttention-3算法正式發布，相關論文《FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision》同期公開。算法初期通過主要作者Tri Dao的個人部落格發布，GitHub上發布了用於Beta測試的原始碼。該算法將開源，並計畫集成到PyTorch和Hugging Face等主流深度學習框架中。此次發布由Colfax Research、Meta、英偉達、Georgia Tech、普林斯頓大學和Together AI等機構的研究團隊聯合完成。
