大模型訓練推理神作，又更新了！

主流大模型都在用的FlashAttention，剛剛升級第三代。

時隔一年，FlashAttention-3已經全方位升級。

訓練速度提升1.5-2倍，FP16下計算吞吐量高達740TFLOPs/s，達理論最大吞吐量75%，更充分利用計算資源，此前只能做到35%。

同時誤差也進一步減小，FP8下的誤差比標準Attention減少2.6倍。

而且這一次，不再是一作Tri Dao單打獨鬥，FlashAttention-3直接和輝達、Meta、Google等合作，針對最強晶片H100專門做最佳化。

輝達CUTLASS團隊和cuDNN團隊，都直接為該研究提供支援。

同時和前作一樣，FlashAttention-3也將開源，PyTorch和Hugging Face中都整合。

曾經在FA2發佈時，我就說過這句話。今天，我想再說一次：

看到CUTLASS和CuTe被用來開讓Tensor Core大顯身手的新演算法，真的泰褲辣。

前Stable Diffusion老闆Emad也非常關注這一進展，他推測使用FlashAttention-3，能將4090的FP8計算吞吐量推升到700+TFLOPs。

自初代發佈以來，FlashAttention已經使大模型速度提高了4-8倍，但還有一個遺憾：尚未充分利用現代 GPU。

針對輝達H100倍後的Hopper架構新特性，三代進行了專門最佳化。

整個系列的核心思路，是IO感知最佳化和分塊處理。

作者認為，傳統的注意力機制效率低的原因，在處理長序列時，會出現記憶體訪問操作頻繁，以及演算法複雜度指數級暴增這兩大問題。

FlashAttention通過IO感知最佳化將資料從較大但緩慢的高頻寬記憶體（HBM）載入到較小但更快的片上記憶體（SRAM），在SRAM中執行計算，減少了記憶體讀寫操作的次數。

分塊處理則是將輸入序列分成若干小塊，每次只處理一個小塊的資料。這種方法使得每次處理的資料量減少，從而降低了記憶體使用和計算複雜度。

這樣一來，兩個關鍵問題就得到瞭解決，這兩大核心思想也在本次的FlashAttention-3中得到了繼承。

但是，第一代的FlashAttention也遺留下了平行性不夠強、工作分區劃分不合理，以及非矩陣乘法較多（GPU計算單元處理矩陣乘法比非矩陣速度更快）的問題。

針對這一問題，第二代FlashAttention通過重寫softmax，減少了重新縮放操作、邊界檢查和因果遮蔽操作的次數，使得大部分計算集中在矩陣乘法上。

另外，FlashAttention-2引入了序列長度維度上的平行化，並針對工作線上程塊之間的分配進行了最佳化，GPU利用效率更高了。

可以說前兩代當中，作者一直堅持著充分利用硬體特點這一思路，但站在今天的視角來看，對硬體的挖掘仍然不夠充分。

到了這次的FlashAttention-3，由於是直接和輝達官方合作，對輝達Hopper架構特點的理解更加透徹，軟硬體之間的協同進一步增強了。

FlashAttention-3的技術報告顯示，為了充分匹配Hopper架構，團隊主要做了三方面的技術升級。

首先，Hopper架構的一個重要特點是Tensor Core的非同步性，FlashAttention-3針對性地提出了一種非同步方式。

具體來說，FlashAttention-3引入了一種“生產者（Producer）-消費者（Consumer）”的程式設計模型，將注意力的計算劃分為兩個角色。

“生產者”負責將資料從HBM非同步載入到片上共用記憶體（SMEM）。這個過程主要利用了Hopper GPU的張量記憶體加速器（TMA），可以在不阻塞CUDA核心的情況下進行資料傳輸。

消費者直接從共用記憶體讀取資料，並使用Tensor Core執行矩陣乘法等計算密集型任務。由於共用記憶體的訪問延遲遠低於全域記憶體，消費者可以快速獲取所需資料，提升計算效率。

為了實現角色的劃分，作者引入了warp專門化技術，用不同的warp分別匹配生產者和消費者，讓兩者可以平行執行。

這其中利用了Hopper架構的動態warp暫存器分配特性，通過setmaxnreg指令最佳化了暫存器資源的利用。

為了進一步提高GPU的利用率，作者又提出了一種“乒乓調度”策略，讓一個warp組執行矩陣乘法時，另一個warp組執行softmax，從而實現計算的重疊。

具體講，FlashAttention-3使用CUDA的同步原語控制不同warp組之間的執行順序，讓不同warp組分別執行兩種運算，然後像乒乓球一樣交替運行。

第二大技術特點，是warp組內部GEMMs和softmax的重疊，核心奧義是重新安排計算的執行順序以提高GPU利用率。

與乒乓調度不同，這裡的計算重排處理的是warp組內部的重疊，而乒乓調度更關注組間協調。

實現方式上，FlashAttention-3提出了一種兩階段GEMM-softmax流水線方案，以打破不同操作之間的資料依賴。

第一階段，當前迭代（iteration）的softmax操作與下一個迭代的Q·K^T矩陣乘法重疊執行。

第二階段，當前迭代的P·V矩陣乘法與下一個迭代的softmax操作重疊執行。

通過引入額外的暫存器和共用記憶體緩衝區，FlashAttention-3實現了跨迭代的資料傳遞和重用。

在每個迭代中，Q·K^T的結果首先儲存在名為S_cur的緩衝區中，用於當前迭代的softmax計算，同時非同步執行下一個迭代的Q·K^T矩陣乘法，結果儲存在名為S_next的緩衝區中。

在執行當前迭代的P·V矩陣乘法時,非同步執行下一個迭代的softmax操作，並更新S_cur和S_next緩衝區。

第三項更新，是用更低的FP8精度替代FP16。

實際上，降低數值精度是一種常見的最佳化策略，可以顯著提高GPU的計算吞吐量和能效，Hopper GPU也引入了FP8精度的Tensor Core支援。

但是，直接將注意力計算從FP16轉換為FP8可能會引入較大的精度損失。

另外，FP8 Tensor Core對輸入資料的佈局也有特定的要求（K維度連續），不幸的是，注意力計算中的輸入資料儲存格式（頭維度連續）並不符合這樣的要求。

所以FlashAttention-3首先引入了一系列記憶體佈局轉換技術，動態轉置V矩陣的塊，改變其連續方式，從而適配FP8 Tensor Core的佈局要求。

在此基礎之上，為了獲得更高的計算精度，FlashAttention-3又採用了分塊量化和非相干處理技術。

傳統的量化方法通常對整個矩陣使用一個統一的縮放因子（per-tensor quantization），無法很好地適應不同區域的數值範圍。

FlashAttention-3則採用了分塊量化（block-wise quantization）的策略，為每個塊單獨設定縮放因子，更好地捕捉局部的數值分佈。

非相干處理（incoherent processing）技術則是通過隨機正交矩陣對輸入資料進行旋轉，破壞不同塊之間的相乾性，減少量化誤差的傳播。

這兩項技術的結合使得FlashAttention-3在FP8精度下取得了更高的計算精度，顯著優於傳統的量化方法。

結果，與基於傳統量化方法的FP8實現相比，FlashAttention-3的使得精度提高了2.6倍。

以上就是FlashAttention-3在充分研究Hopper架構特點後做出的三大更新，針對更新後的表現，作者主要進行了3方面測試。

通過改變序列長度（512、1k、……16k），並設定批大小以確保總token數為16k。研究人員將隱藏維度設定為2048，頭維度設定為64、128或258，計算前向傳播、後向傳播。

對比標準Attention、FlashAttention-2、Triton、cuDNN和FlashAttention-3，在H100 80GB SXM5上FP16的執行階段間。

FlashAttention-3的前向傳播比FlashAttention-2快1.5-2倍，後向傳播快1.5-1.75倍。

與標準Attention相比，FlashAttention-3的速度快了3-16倍。

對於中長序列（1k以上），FlashAttention-3甚至超過了專門為H100最佳化的cuDNN。

在消融實驗中，通過對非因果FP16 FlashAttention-3進行了2階段WGMMA-softmax流水線和warp特殊化的消融研究，參數固定為{batch, seqlen, nheads, hdim} = {4, 8448, 16, 128}。

結果證實，FlashAttention-3改進帶來了顯著加速，從570提升到661。

另外，因為對FlashAttention的數值誤差感興趣，研究團隊還將FlashAttention-2、FlashAttention-3和標準Attention進行了比較。

為了模擬LLMs中的異常特徵和啟動，研究團隊生成了Q、K、V的條目，分佈為：N(0,1)+N(0,100)⋅Bernoulli(0.001)

#FlashAttention三代 你可以在這裡找到鉅亨號中所有使用「FlashAttention三代」為關鍵字所發布的內容