斯坦福博士一己之力讓Attention提速9倍！

【導(dǎo)讀】FlashAttention新升級(jí)！斯坦福博士一人重寫(xiě)算法，第二代實(shí)現(xiàn)了最高9倍速提升。

繼超快且省內(nèi)存的注意力算法FlashAttention爆火后，升級(jí)版的2代來(lái)了。

FlashAttention-2是一種從頭編寫(xiě)的算法，可以加快注意力并減少其內(nèi)存占用，且沒(méi)有任何近似值。

比起第一代，F(xiàn)lashAttention-2速度提升了2倍。

甚至，相較于PyTorch的標(biāo)準(zhǔn)注意力，其運(yùn)行速度最高可達(dá)9倍。

一年前，StanfordAILab博士Tri Dao發(fā)布了FlashAttention，讓注意力快了2到4倍，如今，F(xiàn)lashAttention已經(jīng)被許多企業(yè)和研究室采用，廣泛應(yīng)用于大多數(shù)LLM庫(kù)。

如今，隨著長(zhǎng)文檔查詢(xún)、編寫(xiě)故事等新用例的需要，大語(yǔ)言模型的上下文以前比過(guò)去變長(zhǎng)了許多——GPT-4的上下文長(zhǎng)度是32k，MosaicML的MPT上下文長(zhǎng)度是65k，Anthropic的Claude上下文長(zhǎng)度是100k。

但是，擴(kuò)大Transformer的上下文長(zhǎng)度是一項(xiàng)極大的挑戰(zhàn)，因?yàn)樽鳛槠浜诵牡淖⒁饬拥倪\(yùn)行時(shí)間和內(nèi)存要求，是輸入序列長(zhǎng)度的二次方。

Tri Dao一直在研究FlashAttention-2，它比v1快2倍，比標(biāo)準(zhǔn)的注意力快5到9倍，在A100上已經(jīng)達(dá)到了225 TFLOP/s的訓(xùn)練速度！

論文地址：https://tridao.me/publications/flash2/flash2.pdf

項(xiàng)目地址：https://github.com/Dao-AILab/flash-attention

FlashAttention-2：更好的算法、并行性和工作分區(qū)

端到端訓(xùn)練GPT模型，速度高達(dá)225 TFLOP/s

雖說(shuō)FlashAttention在發(fā)布時(shí)就已經(jīng)比優(yōu)化的基線快了2-4倍，但還是有相當(dāng)大的進(jìn)步空間。

比方說(shuō)，F(xiàn)lashAttention仍然不如優(yōu)化矩陣乘法（GEMM）運(yùn)算快，僅能達(dá)到理論最大FLOPs/s的25-40%（例如，在A100 GPU上的速度可達(dá)124 TFLOPs/s）。

GEMM如何用于卷積

在過(guò)去的幾個(gè)月里，研究人員一直在開(kāi)發(fā)FlashAttention-2，它的性能指標(biāo)比第一代更強(qiáng)。

研究人員表示，2代相當(dāng)于完全從頭重寫(xiě)，使用英偉達(dá)的CUTLASS 3.x及其核心庫(kù)CuTe。從速度上看，F(xiàn)lashAttention-2比之前的版本快了2倍，在A100 GPU上的速度可達(dá)230 TFLOPs/s。

當(dāng)使用端到端來(lái)訓(xùn)練GPT之類(lèi)的語(yǔ)言模型時(shí)，研究人員的訓(xùn)練速度高達(dá)225 TFLOPs/s（模型的FLOP利用率為72%）。

對(duì)注意力計(jì)算重新排序

我們知道，F(xiàn)lashAttention是一種對(duì)注意力計(jì)算進(jìn)行重新排序的算法，利用平鋪、重新計(jì)算來(lái)顯著加快計(jì)算速度，并將序列長(zhǎng)度的內(nèi)存使用量從二次減少到線性。

研究人員將輸入塊從HBM（GPU內(nèi)存）加載到SRAM（快速緩存），并對(duì)該模塊執(zhí)行注意，更新HBM中的輸出。

由于沒(méi)有將大型中間注意力矩陣寫(xiě)入HBM，內(nèi)存的讀/寫(xiě)量也跟著減少，進(jìn)而帶來(lái)了2-4倍的執(zhí)行時(shí)間加速。

下圖是FlashAttention的前向傳遞圖：通過(guò)平鋪和softmax重新縮放，研究人員人員按模塊進(jìn)行操作，避免從HBM讀取或是寫(xiě)入，同時(shí)獲得正確輸出，無(wú)需近似。

然而，F(xiàn)lashAttention仍然存在一些低效率的問(wèn)題，這是由于不同線程塊之間的工作劃分并不理想，以及GPU上的warp——導(dǎo)致低占用率或不必要的共享內(nèi)存讀寫(xiě)。

更少的non-matmul FLOP（非矩陣乘法浮點(diǎn)計(jì)算數(shù)）

研究人員通過(guò)調(diào)整FlashAttention的算法來(lái)減少non-matmul FLOP的次數(shù)。這非常重要，因?yàn)楝F(xiàn)代GPU有專(zhuān)門(mén)的計(jì)算單元（比如英偉達(dá)GPU上的張量核心），這就使得matmul的速度更快。

例如，A100 GPU FP16/BF16 matmul的最大理論吞吐量為312 TFLOPs/s，但non-matmul FP32的理論吞吐量?jī)H為 19.5 TFLOPs/s。

另外，每個(gè)非matmul FLOP比matmul FLOP要貴16倍。

所以為了保持高吞吐量，研究人員希望在matmul FLOP上花盡可能多的時(shí)間。

研究人員還重新編寫(xiě)了FlashAttention中使用的在線softmax技巧，以減少重新縮放操作的數(shù)量，以及邊界檢查和因果掩碼操作，而無(wú)需更改輸出。

更好的并行性

FlashAttention v1在批大小和部數(shù)量上進(jìn)行并行化處理。研究人員使用1個(gè)線程塊來(lái)處理一個(gè)注意力頭，共有 （batch_size * head number） 個(gè)線程塊。

在前向處理（左圖）中，研究者將Worker（線程塊）并行化，每個(gè)Worker負(fù)責(zé)處理注意力矩陣的一個(gè)行塊。在后向處理過(guò)程中（右圖），每個(gè)Worker處理注意力矩陣的一個(gè)列塊

每個(gè)線程塊都在流式多處理器 （SM）運(yùn)行，例如，A100 GPU上有108個(gè)這樣的處理器。當(dāng)這個(gè)數(shù)字很大（比如 ≥80）時(shí)，這種調(diào)度是有效的，因?yàn)樵谶@種情況下，可以有效地使用GPU上幾乎所有的計(jì)算資源。

在長(zhǎng)序列的情況下（通常意味著更小批或更少的頭），為了更好地利用GPU上的多處理器，研究人員在序列長(zhǎng)度的維度上另外進(jìn)行了并行化，使得該機(jī)制獲得了顯著加速。

更好的工作分區(qū)

即使在每個(gè)線程塊內(nèi)，研究人員也必須決定如何在不同的warp（線程束）之間劃分工作（一組32個(gè)線程一起工作）。研究人員通常在每個(gè)線程塊使用4或8個(gè)warp，分區(qū)方案如下圖所示。

研究人員在FlashAttention-2中改進(jìn)了這種分區(qū)，減少了不同warp之間的同步和通信量，從而減少共享內(nèi)存讀/寫(xiě)。

對(duì)于每個(gè)塊，F(xiàn)lashAttention將K和V分割到4個(gè)warp上，同時(shí)保持Q可被所有warp訪問(wèn)。這稱(chēng)為「sliced-K」方案。

然而，這樣做的效率并不高，因?yàn)樗衱arp都需要將其中間結(jié)果寫(xiě)入共享內(nèi)存，進(jìn)行同步，然后再將中間結(jié)果相加。

而這些共享內(nèi)存讀/寫(xiě)會(huì)減慢FlashAttention中的前向傳播速度。

在FlashAttention-2中，研究人員將Q拆分為4個(gè)warp，同時(shí)保持所有warp都可以訪問(wèn)K和V。

在每個(gè)warp執(zhí)行矩陣乘法得到Q K^T的一個(gè)切片后，它們只需與共享的V切片相乘，即可得到相應(yīng)的輸出切片。

這樣一來(lái)，warp之間就不再需要通信。共享內(nèi)存讀寫(xiě)的減少就可以提高速度。

新功能：頭的維度高達(dá)256，多查詢(xún)注意力

FlashAttention僅支持最大128的頭的維度，雖說(shuō)適用于大多數(shù)模型，但還是有一些模型被排除在外。

FlashAttention-2現(xiàn)在支持256的頭的維度，這意味著GPT-J、CodeGen、CodeGen2以及Stable Diffusion 1.x等模型都可以使用FlashAttention-2來(lái)獲得加速和節(jié)省內(nèi)存。

v2還支持多查詢(xún)注意力（MQA）以及分組查詢(xún)注意力（GQA）。

GQA為每組查詢(xún)頭共享單個(gè)key和value的頭，在多頭和多查詢(xún)注意之間進(jìn)行插值

這些都是注意力的變體，其中多個(gè)查詢(xún)頭會(huì)指向key和value的同一個(gè)頭，以減少推理過(guò)程中KV緩存的大小，并可以顯著提高推理的吞吐量。

注意力基準(zhǔn)

研究人員人員在A100 80GB SXM4 GPU 上測(cè)量不同設(shè)置（有無(wú)因果掩碼、頭的維度是64或128）下不同注意力方法的運(yùn)行時(shí)間。

研究人員發(fā)現(xiàn)FlashAttention-2比第一代快大約2倍（包括在xformers庫(kù)和Triton中的其他實(shí)現(xiàn)）。

與PyTorch中的標(biāo)準(zhǔn)注意力實(shí)現(xiàn)相比，F(xiàn)lashAttention-2的速度最高可達(dá)其9倍。

A100 GPU上的前向+后向速度

只需在H100 GPU上運(yùn)行相同的實(shí)現(xiàn)（不需要使用特殊指令來(lái)利用TMA和第四代Tensor Core等新硬件功能），研究人員就可以獲得高達(dá)335 TFLOPs/s的速度。

H100 GPU上的前向+后向速度

當(dāng)用于端到端訓(xùn)練GPT類(lèi)模型時(shí)，F(xiàn)lashAttention-2能在A100 GPU上實(shí)現(xiàn)高達(dá)225TFLOPs/s的速度（模型FLOPs利用率為72%）。

與已經(jīng)非常優(yōu)化的FlashAttention模型相比，端到端的加速進(jìn)一步提高了1.3倍。

未來(lái)的工作

速度上快2倍，意味著研究人員可以用與之前訓(xùn)練8k上下文模型相同的成本，來(lái)訓(xùn)練16k上下文長(zhǎng)度的模型。這些模型可以理解長(zhǎng)篇書(shū)籍和報(bào)告、高分辨率圖像、音頻和視頻。

同時(shí)，F(xiàn)lashAttention-2還將加速現(xiàn)有模型的訓(xùn)練、微調(diào)和推理。

在不久的將來(lái)，研究人員還計(jì)劃擴(kuò)大合作，使FlashAttention廣泛適用于不同類(lèi)型的設(shè)備（例如H100 GPU、AMD GPU）以及新的數(shù)據(jù)類(lèi)型（例如fp8）。

下一步，研究人員計(jì)劃針對(duì)H100 GPU進(jìn)一步優(yōu)化FlashAttention-2，以使用新的硬件功能（TMA、第四代Tensor Core、fp8等等）。

將FlashAttention-2中的低級(jí)優(yōu)化與高級(jí)算法更改（例如局部、擴(kuò)張、塊稀疏注意力）相結(jié)合，可以讓研究人員用更長(zhǎng)的上下文來(lái)訓(xùn)練AI模型。

研究人員也很高興與編譯器研究人員合作，使這些優(yōu)化技術(shù)更好地應(yīng)用于編程。