H100利用率飆升至75%！英偉達(dá)親自下場(chǎng)FlashAttention三代升級(jí)，比標(biāo)準(zhǔn)注意力快16倍

作者：量子位 2024-07-12 11:35:49

人工智能新聞

訓(xùn)練速度提升1.5-2倍，F(xiàn)P16下計(jì)算吞吐量高達(dá)740TFLOPs/s，達(dá)理論最大吞吐量75%，更充分利用計(jì)算資源，此前只能做到35%。

大模型訓(xùn)練推理神作，又更新了！

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

時(shí)隔一年，F(xiàn)lashAttention-3已經(jīng)全方位升級(jí)。

訓(xùn)練速度提升1.5-2倍，F(xiàn)P16下計(jì)算吞吐量高達(dá)740TFLOPs/s，達(dá)理論最大吞吐量75%，更充分利用計(jì)算資源，此前只能做到35%。

FP8下速度接近1.2PFLOPs/s！

同時(shí)誤差也進(jìn)一步減小，F(xiàn)P8下的誤差比標(biāo)準(zhǔn)Attention減少2.6倍。

而且這一次，不再是一作Tri Dao單打獨(dú)斗，F(xiàn)lashAttention-3直接和英偉達(dá)、Meta、谷歌等合作，針對(duì)最強(qiáng)芯片H100專門做優(yōu)化。

英偉達(dá)CUTLASS團(tuán)隊(duì)和cuDNN團(tuán)隊(duì)，都直接為該研究提供支持。

同時(shí)和前作一樣，F(xiàn)lashAttention-3也將開源，PyTorch和Hugging Face中都集成。

作者之一Vijay Thakkar激動(dòng)表示：

曾經(jīng)在FA2發(fā)布時(shí)，我就說(shuō)過(guò)這句話。今天，我想再說(shuō)一次：
看到CUTLASS和CuTe被用來(lái)開讓Tensor Core大顯身手的新算法，真的泰褲辣。

前Stable Diffusion老板Emad也非常關(guān)注這一進(jìn)展，他推測(cè)使用FlashAttention-3，能將4090的FP8計(jì)算吞吐量推升到700+TFLOPs。

充分利用Hopper架構(gòu)特點(diǎn)

自初代發(fā)布以來(lái)，F(xiàn)lashAttention已經(jīng)使大模型速度提高了4-8倍，但還有一個(gè)遺憾：尚未充分利用現(xiàn)代 GPU。

針對(duì)英偉達(dá)H100倍后的Hopper架構(gòu)新特性，三代進(jìn)行了專門優(yōu)化。

整個(gè)系列的核心思路，是IO感知優(yōu)化和分塊處理。

作者認(rèn)為，傳統(tǒng)的注意力機(jī)制效率低的原因，在處理長(zhǎng)序列時(shí)，會(huì)出現(xiàn)內(nèi)存訪問(wèn)操作頻繁，以及算法復(fù)雜度指數(shù)級(jí)暴增這兩大問(wèn)題。

FlashAttention通過(guò)IO感知優(yōu)化將數(shù)據(jù)從較大但緩慢的高帶寬內(nèi)存（HBM）加載到較小但更快的片上內(nèi)存（SRAM），在SRAM中執(zhí)行計(jì)算，減少了內(nèi)存讀寫操作的次數(shù)。

分塊處理則是將輸入序列分成若干小塊，每次只處理一個(gè)小塊的數(shù)據(jù)。這種方法使得每次處理的數(shù)據(jù)量減少，從而降低了內(nèi)存使用和計(jì)算復(fù)雜度。

這樣一來(lái)，兩個(gè)關(guān)鍵問(wèn)題就得到了解決，這兩大核心思想也在本次的FlashAttention-3中得到了繼承。

△第一代FlashAttention原理圖

但是，第一代的FlashAttention也遺留下了并行性不夠強(qiáng)、工作分區(qū)劃分不合理，以及非矩陣乘法較多（GPU計(jì)算單元處理矩陣乘法比非矩陣速度更快）的問(wèn)題。

針對(duì)這一問(wèn)題，第二代FlashAttention通過(guò)重寫softmax，減少了重新縮放操作、邊界檢查和因果屏蔽操作的次數(shù)，使得大部分計(jì)算集中在矩陣乘法上。

另外，F(xiàn)lashAttention-2引入了序列長(zhǎng)度維度上的并行化，并針對(duì)工作在線程塊之間的分配進(jìn)行了優(yōu)化，GPU利用效率更高了。

可以說(shuō)前兩代當(dāng)中，作者一直堅(jiān)持著充分利用硬件特點(diǎn)這一思路，但站在今天的視角來(lái)看，對(duì)硬件的挖掘仍然不夠充分。

到了這次的FlashAttention-3，由于是直接和英偉達(dá)官方合作，對(duì)英偉達(dá)Hopper架構(gòu)特點(diǎn)的理解更加透徹，軟硬件之間的協(xié)同進(jìn)一步增強(qiáng)了。

FlashAttention-3的技術(shù)報(bào)告顯示，為了充分匹配Hopper架構(gòu)，團(tuán)隊(duì)主要做了三方面的技術(shù)升級(jí)。

首先，Hopper架構(gòu)的一個(gè)重要特點(diǎn)是Tensor Core的異步性，F(xiàn)lashAttention-3針對(duì)性地提出了一種異步方式。

具體來(lái)說(shuō)，F(xiàn)lashAttention-3引入了一種“生產(chǎn)者（Producer）-消費(fèi)者（Consumer）”的編程模型，將注意力的計(jì)算劃分為兩個(gè)角色。

“生產(chǎn)者”負(fù)責(zé)將數(shù)據(jù)從HBM異步加載到片上共享內(nèi)存（SMEM）。這個(gè)過(guò)程主要利用了Hopper GPU的張量?jī)?nèi)存加速器（TMA），可以在不阻塞CUDA核心的情況下進(jìn)行數(shù)據(jù)傳輸。
消費(fèi)者直接從共享內(nèi)存讀取數(shù)據(jù)，并使用Tensor Core執(zhí)行矩陣乘法等計(jì)算密集型任務(wù)。由于共享內(nèi)存的訪問(wèn)延遲遠(yuǎn)低于全局內(nèi)存，消費(fèi)者可以快速獲取所需數(shù)據(jù)，提升計(jì)算效率。

△Hopper中的張量?jī)?nèi)存加速器

為了實(shí)現(xiàn)角色的劃分，作者引入了warp專門化技術(shù)，用不同的warp分別匹配生產(chǎn)者和消費(fèi)者，讓兩者可以并行執(zhí)行。

這其中利用了Hopper架構(gòu)的動(dòng)態(tài)warp寄存器分配特性，通過(guò)setmaxnreg指令優(yōu)化了寄存器資源的利用。

為了進(jìn)一步提高GPU的利用率，作者又提出了一種“乒乓調(diào)度”策略，讓一個(gè)warp組執(zhí)行矩陣乘法時(shí)，另一個(gè)warp組執(zhí)行softmax，從而實(shí)現(xiàn)計(jì)算的重疊。

具體講，F(xiàn)lashAttention-3使用CUDA的同步原語(yǔ)控制不同warp組之間的執(zhí)行順序，讓不同warp組分別執(zhí)行兩種運(yùn)算，然后像乒乓球一樣交替運(yùn)行。

第二大技術(shù)特點(diǎn)，是warp組內(nèi)部GEMMs和softmax的重疊，核心奧義是重新安排計(jì)算的執(zhí)行順序以提高GPU利用率。

與乒乓調(diào)度不同，這里的計(jì)算重排處理的是warp組內(nèi)部的重疊，而乒乓調(diào)度更關(guān)注組間協(xié)調(diào)。

實(shí)現(xiàn)方式上，F(xiàn)lashAttention-3提出了一種兩階段GEMM-softmax流水線方案，以打破不同操作之間的數(shù)據(jù)依賴。

第一階段，當(dāng)前迭代（iteration）的softmax操作與下一個(gè)迭代的Q·K^T矩陣乘法重疊執(zhí)行。
第二階段，當(dāng)前迭代的P·V矩陣乘法與下一個(gè)迭代的softmax操作重疊執(zhí)行。

通過(guò)引入額外的寄存器和共享內(nèi)存緩沖區(qū)，F(xiàn)lashAttention-3實(shí)現(xiàn)了跨迭代的數(shù)據(jù)傳遞和重用。

在每個(gè)迭代中，Q·K^T的結(jié)果首先存儲(chǔ)在名為S_cur的緩沖區(qū)中，用于當(dāng)前迭代的softmax計(jì)算，同時(shí)異步執(zhí)行下一個(gè)迭代的Q·K^T矩陣乘法，結(jié)果存儲(chǔ)在名為S_next的緩沖區(qū)中。

在執(zhí)行當(dāng)前迭代的P·V矩陣乘法時(shí),異步執(zhí)行下一個(gè)迭代的softmax操作，并更新S_cur和S_next緩沖區(qū)。

第三項(xiàng)更新，是用更低的FP8精度替代FP16。

實(shí)際上，降低數(shù)值精度是一種常見的優(yōu)化策略，可以顯著提高GPU的計(jì)算吞吐量和能效，Hopper GPU也引入了FP8精度的Tensor Core支持。

但是，直接將注意力計(jì)算從FP16轉(zhuǎn)換為FP8可能會(huì)引入較大的精度損失。

另外，F(xiàn)P8 Tensor Core對(duì)輸入數(shù)據(jù)的布局也有特定的要求（K維度連續(xù)），不幸的是，注意力計(jì)算中的輸入數(shù)據(jù)存儲(chǔ)格式（頭維度連續(xù)）并不符合這樣的要求。

所以FlashAttention-3首先引入了一系列內(nèi)存布局轉(zhuǎn)換技術(shù)，動(dòng)態(tài)轉(zhuǎn)置V矩陣的塊，改變其連續(xù)方式，從而適配FP8 Tensor Core的布局要求。

在此基礎(chǔ)之上，為了獲得更高的計(jì)算精度，F(xiàn)lashAttention-3又采用了分塊量化和非相干處理技術(shù)。

傳統(tǒng)的量化方法通常對(duì)整個(gè)矩陣使用一個(gè)統(tǒng)一的縮放因子（per-tensor quantization），無(wú)法很好地適應(yīng)不同區(qū)域的數(shù)值范圍。

FlashAttention-3則采用了分塊量化（block-wise quantization）的策略，為每個(gè)塊單獨(dú)設(shè)置縮放因子，更好地捕捉局部的數(shù)值分布。

非相干處理（incoherent processing）技術(shù)則是通過(guò)隨機(jī)正交矩陣對(duì)輸入數(shù)據(jù)進(jìn)行旋轉(zhuǎn)，破壞不同塊之間的相干性，減少量化誤差的傳播。

這兩項(xiàng)技術(shù)的結(jié)合使得FlashAttention-3在FP8精度下取得了更高的計(jì)算精度，顯著優(yōu)于傳統(tǒng)的量化方法。

結(jié)果，與基于傳統(tǒng)量化方法的FP8實(shí)現(xiàn)相比，F(xiàn)lashAttention-3的使得精度提高了2.6倍。

比標(biāo)準(zhǔn)Attention快16倍

以上就是FlashAttention-3在充分研究Hopper架構(gòu)特點(diǎn)后做出的三大更新，針對(duì)更新后的表現(xiàn)，作者主要進(jìn)行了3方面測(cè)試。

注意力基準(zhǔn)測(cè)試
消融實(shí)驗(yàn)
FP8注意力準(zhǔn)確性測(cè)試

首先來(lái)看注意力基準(zhǔn)測(cè)試。

通過(guò)改變序列長(zhǎng)度（512、1k、……16k），并設(shè)置批大小以確保總token數(shù)為16k。研究人員將隱藏維度設(shè)置為2048，頭維度設(shè)置為64、128或258，計(jì)算前向傳播、后向傳播。

對(duì)比標(biāo)準(zhǔn)Attention、FlashAttention-2、Triton、cuDNN和FlashAttention-3，在H100 80GB SXM5上FP16的運(yùn)行時(shí)間。

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

與標(biāo)準(zhǔn)Attention相比，F(xiàn)lashAttention-3的速度快了3-16倍。

對(duì)于中長(zhǎng)序列（1k以上），F(xiàn)lashAttention-3甚至超過(guò)了專門為H100優(yōu)化的cuDNN。

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

結(jié)果證實(shí)，F(xiàn)lashAttention-3改進(jìn)帶來(lái)了顯著加速，從570提升到661。