FlashAttention-4震撼來襲，原生支持Blackwell GPU，英偉達的護城河更深了？

（來源：機器之心）

在正在舉辦的半導體行業會議 Hot Chips 2025 上，TogetherAI 首席科學家 Tri Dao 公佈了 FlashAttention-4。

據介紹，在 Backwell 上，FlashAttention-4 的速度比英偉達 cuDNN 庫中的注意力覈實現快可達 22%！

在這個新版本的 FlashAttention 中，Tri Dao 團隊實現了兩項關鍵的算法改進。

一、它使用了一種新的在線 softmax 算法，可跳過了 90% 的輸出 rescaling。

二、爲了更好地將 softmax 計算與張量覈計算重疊，它使用了指數 (MUFU.EX2) 的軟件模擬來提高吞吐量。

此外，FlashAttention-4 使用的是 CUTLASS CuTe Python DSL，其移植到 ROCm HIP 的難度要高出 10 倍，而 CUDA C++ 移植到 ROCm HIP 則更容易。

有意思的是，Tri Dao 還宣佈，在執行 A@B+C 計算時，對於 Blackwell 上在歸約維度 K 較小的計算場景中，他使用 CUTLASS CuTe-DSL 編寫的核（kernel）比英偉達最新的 cuBLAS 13.0 庫快不少。而在標準矩陣算法 A@B 時，兩者速度總體是相當的。

據介紹，他的核通過使用兩個累積緩衝區來重疊 epilogue，從而擊敗了 cuBLAS。

Semi Analysis 表示，像 Tri Dao 這樣的開發者是 CUDA 護城河的核心優勢之一，因為 Tri Dao 只使用英偉達 GPU，並將其大部分核開源給其他英偉達開發者羣體。Tri Dao 等研究者均不使用 ROCm AMD GPU 或 Trainium 芯片。

這對於 AMD 等來説可不是好消息，假如 AMD 希望 Tri Dao 和他的團隊在 ROCm 上實現算法突破。那麼，它就應該為 TogetherAI GPU 雲服務上的 AMD GPU 提供優惠支持。Semi Analysis 分析説：「谷歌為 Noam Shazeer 支付了 27 億美元，Zucc 為 OpenAI 工程師支付了 1 億美元，AMD 擁有足夠的現金，可以為 TogetherAI/Tri Dao 支付 5000 萬美元來啟動 ROCm 生態系統。」

FlashAttention 最早由 Tri Dao 等人在 2022 年提出，論文標題為《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》。

論文地址：https://arxiv.org/pdf/2205.14135

其背景是傳統的注意力機制因需生成 N×N 的注意力矩陣，在序列長度 N 增長時引發二次的（quadratic）時間和內存開銷。

而 FlashAttention 強調「IO-awareness」，不再將注意力矩陣完整載入，而是通過「tiling+softmax rescaling」策略，將數據塊臨時存入高速緩存（SRAM），在內部積累，再寫回高帶寬內存（HBM），避免了大量讀寫開銷，內存複雜度得到顯著降低 —— 從 O (N²) 降至 O (N)。

如圖所示，在左圖中，FlashAttention 使用了 tiling 技術來防止在（相對較慢的）GPU HBM 上執行很大的 𝑁 × 𝑁 注意力矩陣（虛線框）。在外層循環（紅色箭頭）中，FlashAttention 循環遍歷 K 和 V 矩陣的塊，並將其加載到快速片上 SRAM 中。在每個塊中，FlashAttention 循環遍歷 Q 矩陣的塊（藍色箭頭），將其加載到 SRAM 中，並將注意力計算的輸出寫回 HBM。

在右圖中，可以看到相比 GPT-2 上 PyTorch 注意力實現，FlashAttention 速度更快 ——FlashAttention 無需將大型 𝑁 × 𝑁 注意力矩陣讀寫到 HBM，從而將注意力計算速度提升了 7.6 倍。

整體上，初代 FlashAttention 帶來的增益也很顯著：在 BERT-large（序列長度 512）中相比 MLPerf 基線提升訓練速度約 15%；GPT-2（序列長度 1K）提升約 3 倍；在 Long-Range Arena（序列長度 1K–4K）提升約 2.4 倍。

一年后，FlashAttention-2 問世，這一次，作者僅 Tri Dao 一人。順帶一提，他還在這一年的晚些時候與 Albert Gu 共同提出了 Mamba。

論文地址：https://arxiv.org/pdf/2307.08691

其改進的焦點是：FlashAttention 已顯著提升性能，但在 GPU 上仍存在低吞吐率的問題，僅能達到理論峰值很低的比例（約 25–40%）。

為此，Tri Dao 提出的解決策略包括：

• 工作劃分優化：重新設計分塊策略與線程分配，提升並行效率，增加硬件利用率；

• 減少非矩陣運算，加快整體執行；

• 支持更大 head size（至 256） 及多查詢注意力（MQA） 和分組查詢注意力（GQA），適配更多模型架構需求。

結果，相比初代 FlashAttention，FlashAttention-2 速度提高約 2–4×；在 A100 GPU 上 FP16/BF16 可達到高至 230 TFLOPs/s，達 PyTorch 標準實現 9 倍速度提升。參閱機器之心報道《比標準 Attention 提速 5-9 倍，大模型都在用的 FlashAttention v2 來了》。

又一年，FlashAttention-3 誕生，這一次改進的重點是適配 Hopper 架構，異步與低精度。可以看到，Tri Dao 這一次的名字掛在最后。此時他雖然還繼續在普林斯頓大學任教，但也同時已經是 Together AI 的首席科學家。

論文地址：https://arxiv.org/pdf/2407.08608

爲了能加速在 Hopper GPU 上的注意力，FlashAttention-3 主要採用了三種技術：

• 通過 warp-specialization 重疊整體計算和數據移動；

• 交錯分塊 matmul 和 softmax 運算；

• 利用硬件支持 FP8 低精度的不連貫處理。

FlashAttention-3 的速度是 FlashAttention-2 的 1.5-2.0 倍，高達 740 TFLOPS，即 H100 理論最大 FLOPS 利用率為 75%。使用 FP8，FlashAttention-3 的速度更是接近 1.2 PFLOPS。參閱機器之心報道《英偉達又賺到了！FlashAttention3 來了：H100 利用率飆升至 75%》。

現在，到了 2025 年，FlashAttention-4 準時到來，增加了對 Blackwell GPU 的原生支持——之前，想要在 Blackwell 上跑 FlashAttention，如果直接用開源倉庫，常常會遇到編譯錯誤、kernel 缺失或性能未優化的情況，可用的 Blackwell 加速主要是藉助英偉達 Triton/cuDNN 的間接支持。

圖源：https://www.reddit.com/r/LocalLLaMA/comments/1mt9htu/flashattention_4_leak/

此時，FlashAttention 的 GitHub 軟件庫已經積累了超過 1.91 萬星。

項目地址：https://github.com/Dao-AILab/flash-attention

目前，Tri Dao 團隊尚未發佈 FlashAttention-4 的技術報告，更多細節還有待進一步揭曉。

參考鏈接

https://x.com/tri_dao/status/1960217005446791448

https://x.com/SemiAnalysis_/status/1960070677379133949

https://www.reddit.com/r/LocalLLaMA/comments/1mt9htu/flashattention_4_leak/