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

機(jī)器之心編輯部

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

近日，深度學(xué)習(xí)領(lǐng)域重要底層優(yōu)化技術(shù) FlashAttention 迎來(lái)大版本更新。

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

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

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

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

接下來(lái)，就來(lái)詳細(xì)了解一下。

硬件趨勢(shì)：不對(duì)稱的硬件擴(kuò)展

長(zhǎng)期以來(lái)，Attention 作為無(wú)處不在的 Transformer 架構(gòu)中的核心層，一直是大語(yǔ)言模型和長(zhǎng)上下文應(yīng)用的性能瓶頸。

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

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

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

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

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

具體來(lái)看，Attention 的核心包含兩個(gè)通用矩陣乘法（GEMM）：

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

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

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

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

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

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

Blackwell 的新硬件特性

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

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

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

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

編程語(yǔ)言與框架：CuTe-DSL

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

Kernel 代碼使用 Python 編寫(xiě)，隨后 DSL 會(huì)將其降級(jí)（lower 為 PTX，再由 CUDA 工具鏈編譯為 GPU 機(jī)器代碼。

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

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

Attention 性能基準(zhǔn)測(cè)試

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

結(jié)果顯示：

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

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

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

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

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

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

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

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

更多內(nèi)容，可查看論文原文獲?。?/p>

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

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