FlashAttention-4正式發(fā)布：算法流水線大改，矩陣乘法級速度

機器之心編輯部

經(jīng)過一年的努力，F(xiàn)lashAttention-4 終于正式上線了。

近日，深度學習領域重要底層優(yōu)化技術 FlashAttention 迎來大版本更新。

FlashAttention 核心作者、普林斯頓大學助理教授 Tri Dao 表示，在 Blackwell GPU 上，即使瓶頸截然不同，注意力機制的執(zhí)行速度現(xiàn)在也幾乎與矩陣乘法一樣快了！

當前，Tensor Core 的速度現(xiàn)在非常快，以至于注意力前向傳播的瓶頸呈指數(shù)級增長，而注意力后向傳播的瓶頸是共享內(nèi)存帶寬。

重新設計的算法中包含一些旨在克服這些瓶頸的機制，包括使用多項式進行指數(shù)模擬，新的在線 softmax 可以避免 90% 的 softmax 重新縮放，2CTA MMA 指令允許兩個線程塊共享操作數(shù)以減少 smem 流量等。

• 論文地址：https://github.com/Dao-AILab/flash-attention/blob/main/assets/fa4_paper.pdf
• 代碼鏈接：https://github.com/Dao-AILab/flash-attention

接下來，就來詳細了解一下。

硬件趨勢：不對稱的硬件擴展

長期以來，Attention 作為無處不在的 Transformer 架構中的核心層，一直是大語言模型和長上下文應用的性能瓶頸。

此前 FlashAttention-3 通過異步執(zhí)行和 warp 專門化對 Attention 進行了優(yōu)化，但其主要針對的是 Hopper GPU（H100）架構。

然而，AI 行業(yè)已經(jīng)迅速轉向部署 Blackwell 架構系統(tǒng)，例如 B200 和 GB200。而像 Blackwell GPU 這樣的現(xiàn)代加速器延續(xù)了一種趨勢：硬件的非對稱擴展（asymmetric hardware scaling）。

在這種趨勢下，張量核心（Tensor Core）的吞吐量增長速度遠快于其他硬件資源，像是共享內(nèi)存帶寬、用于指數(shù)運算等超越函數(shù)運算的特殊函數(shù)單元（SFU），以及通用整數(shù)與浮點 ALU……

舉個例子，從 Hopper H100 到 Blackwell B200，BF16 張量核心吞吐量增加了 2.25 倍（從 1 到 2.25PFLOPs），但 SFU 數(shù)量和共享內(nèi)存帶寬基本保持不變。

這種擴展不對稱性對像 Attention 這樣的復雜 kernel 優(yōu)化產(chǎn)生了深遠影響。

具體來看，Attention 的核心包含兩個通用矩陣乘法（GEMM）：

中間夾著 softmax，但在真實實踐中，Attention 還涉及大量輔助工作，比如數(shù)據(jù)搬運、同步、數(shù)據(jù)布局轉換、元素級運算、調度、mask 處理等。

傳統(tǒng)的觀點認為，Attention 的性能完全由 GEMM 的速度決定。然而，對 B200 進行「速度與饋送」分析顯示：主要的瓶頸不在于張量核心，而是：

1. 前向傳播中用于 Softmax 指數(shù)運算的 SFU 單元；
2. 反向傳播中的共享內(nèi)存流量，受 shared memory bandwidth 限制。

為此，團隊推出FlashAttention-4，一種算法 + kernel 的協(xié)同設計，核心目標在于，通過最大化矩陣乘法與其他瓶頸資源之間的重疊，在 B200（BF16）上，最高可達 1605TFLOPs/s（71% 的利用率），比 cuDNN 9.13 快 1.3 倍，比 Triton 快 2.7 倍。

協(xié)同設計的核心思路如下：

• 新型流水線：為前向和反向傳播分別設計了新的軟件流水線，利用 Blackwell 的全異步 MMA 和更大分塊（Tile）尺寸，最大化 Tensor Core 計算、softmax 計算以及內(nèi)存操作之間的重疊執(zhí)行；
• 前向傳播 (FWD)：在 FMA 單元上通過多項式近似實現(xiàn)指數(shù)函數(shù)的軟件仿真，以提升指數(shù)計算吞吐量；同時引入條件式 softmax 重縮放（conditional softmax rescaling），跳過不必要的重縮放操作，從而緩解 SFU 瓶頸；
• 反向傳播 (BWD)：利用張量內(nèi)存 (TMEM) 存儲中間結果，以緩解共享內(nèi)存流量壓力；同時，結合 Blackwell 新增的 2-CTA MMA 模式，進一步降低共享內(nèi)存訪問，并將 atomic reduction 次數(shù)減少一半；此外，還支持確定性執(zhí)行模式，以實現(xiàn)可復現(xiàn)訓練；
• 調度優(yōu)化：引入新的 tile 調度器，解決因果掩碼和變長序列導致的負載不均衡。

Blackwell 的新硬件特性

張量內(nèi)存（TMEM）：在 B200 上，148 個 SM（流式多處理器）中的每一個都配備了 256 KB 的 TMEM，與 Tensor Core 直接連接，用于 warp 同步的中間結果存儲。

完全異步的第五代張量核心：指令 tcgen05.mma 支持異步執(zhí)行，并將累加結果存儲在 TMEM 中。對于 BF16 和 FP16，單個 CTA 可使用的最大 UMMA tile 為 128×256×16，約為 Hopper 架構中最大 WGMMA 原子塊的 2 倍。UMMA 由單個線程發(fā)起，從而減輕寄存器壓力，使得在不出現(xiàn) Hopper warpgroup MMA 那種寄存器溢出問題的情況下，可以更容易地使用更大的 tile 和更深的流水線。

此外，這也使 warp 專門化更具可行性：部分 warp 負責搬運 tile，另一些 warp 負責發(fā)起 MMA，從而實現(xiàn)矩陣乘加運算與 softmax 計算以及內(nèi)存訪問的重疊執(zhí)行。tcgen05.mma 還可以直接從 TMEM 中讀取操作數(shù) A。

2-CTA MMA：Blackwell 支持在同一 cluster 中由一對 CTA 共同執(zhí)行一個 UMMA 運算，并跨越兩個 CTA 的 TMEM。由 leader CTA 中的一個線程發(fā)起 MMA，但在執(zhí)行期間兩個 CTA 都必須保持活躍。通過在這對 CTA 之間拆分 M 和 N 維度，可以將 MMA 的 tile 尺寸擴展到 256×256×16，從而減少冗余數(shù)據(jù)傳輸并降低每個 CTA 的資源占用。在一個 kernel 中，CTA 組大小（1 或 2）在 TMEM 操作和 Tensor Core 運算之間必須保持一致。

編程語言與框架：CuTe-DSL

FlashAttention-4（FA4）完全使用 CuTe-DSL 實現(xiàn)，這是 CUTLASS 提供的 Python kernel DSL。

Kernel 代碼使用 Python 編寫，隨后 DSL 會將其降級（lower 為 PTX，再由 CUDA 工具鏈編譯為 GPU 機器代碼。

該編程模型在抽象層面與 CuTe / CUTLASS 保持一致，同時提供 PTX 級別的 escape hatch（底層控制接口）。與使用 C++ 模板相比，這種方式可以將編譯時間縮短約 20–30 倍。

對此，Tri Dao 更是在 X 上發(fā)帖稱感到「莫名興奮」，這意味著，安裝 /「編譯」現(xiàn)在只需幾秒鐘，而不是幾分鐘 / 幾小時。

Attention 性能基準測試

團隊展示了 FlashAttention-4 在 B200（BF16）上的性能結果，并將其與 FlashAttention-2 以及 Triton、Gluon 和 cuDNN 的實現(xiàn)進行了對比。

結果顯示：

• 前向傳播（forward pass）：FlashAttention-4 比 cuDNN 9.13 快 1.1–1.3 倍，比 Triton 實現(xiàn)快 2.1–2.7 倍。
• 反向傳播（backward pass）：在長序列長度場景下，F(xiàn)lashAttention-4 的表現(xiàn)始終優(yōu)于其他基準模型。

而 FlashAttention-4 一經(jīng)發(fā)布，也引起了大家的熱議。

Pytorch 官方宣布 FlexAttention 現(xiàn)已支持 FlashAttention-4 后端。

Pytorch 表示，很長一段時間以來，F(xiàn)lexAttention 讓研究人員能夠快速原型化各種自定義 Attention 變體，目前已有 1000 多個代碼倉庫采用，并有數(shù)十篇論文對其進行了引用。

然而，用戶常常會遇到性能瓶頸，直到 FlashAttention-4 的出現(xiàn)。

如今，他們已在 Hopper 和 Blackwell GPU 上為 FlexAttention 增加了 FlashAttention-4 后端。PyTorch 現(xiàn)在可以自動生成 CuTeDSL 的 score/mask 修改代碼，并通過 JIT 編譯為自定義 Attention 變體實例化 FlashAttention-4。

結果顯示，在算力受限的工作負載下，相比 Triton，仍可實現(xiàn) 1.2 倍到 3.2 倍的性能提升。研究人員再也不必在「靈活性」和「高性能」之間做單選題。

一位網(wǎng)友則認為，「FlashAttention-4 是一個里程碑。」在 Blackwell 架構上，Attention 已經(jīng)能夠達到接近矩陣乘法（matmul）速度，這意味著計算瓶頸將完全轉移到內(nèi)存與通信上。約 1600TFLOPs 的 Attention 性能堪稱驚人 —— 相比 FlashAttention-3 提升了 2–3 倍。「這將直接惠及所有前沿大模型。」因為，更快的 Attention 意味著更長的有效上下文窗口、更低的推理成本、更強的規(guī)模化推理能力……

更多內(nèi)容，可查看論文原文獲取！

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

https://tridao.me/blog/2026/flash4/