豐色 發(fā)自 凹非寺

|  

這兩天，F(xiàn)lashAttention團(tuán)隊(duì)推出了新作：

一種給Transformer架構(gòu)大模型推理加速的新方法，最高可提速8倍。

該方法尤其造福于長(zhǎng)上下文LLM，在64k長(zhǎng)度的CodeLlama-34B上通過了驗(yàn)證。

甚至得到了PyTorch官方認(rèn)可：

如果你之前有所關(guān)注，就會(huì)記得用FlashAttention給大模型加速效果真的很驚艷。

不過它僅限于訓(xùn)練階段。

因此，這一新成果一出，就有網(wǎng)友表示：

等推理加速等了好久，終于來了。

據(jù)介紹，這個(gè)新方法也是在FlashAttention的基礎(chǔ)之上衍生而出，主要思想也不復(fù)雜：

用并行操作盡快加載Key和Value緩存，然后分別重新縮放再合并結(jié)果，最終獲得推理速度上的大幅提升。
詳細(xì)來看。

提速8倍的長(zhǎng)上下文推理方法來了

該方法被命名為Flash-Decoding。

背景與動(dòng)機(jī)

根據(jù)作者介紹：

LLM的推理（即“解碼”）過程是迭代的，即一次生成一個(gè)token，組成一個(gè)完整句子需要n個(gè)token以及n次前向傳遞。

不過，由于我們可以緩存之前計(jì)算出來的token，所以單個(gè)生成步驟并不總是依賴于上下文長(zhǎng)度。

但有一個(gè)操作例外：注意力?(attention)，它不能隨著上下文長(zhǎng)度靈活擴(kuò)展。

鑒于長(zhǎng)上下文已成趨勢(shì)，比如目前最大的開源LLM已達(dá)100k（CodeLlama），我們不得不注意到attention在大模型推理過程中浪費(fèi)了太多時(shí)間，時(shí)間就是金錢。

更別提attention在batch size上進(jìn)行擴(kuò)展時(shí)，即使模型上下文相對(duì)較短，它也可能成為性能瓶頸（因?yàn)槟Ｐ鸵x取的內(nèi)存量與batch size成比例，而它僅取決于模型其余部分的大小）。

怎么破解？

不可復(fù)用的FlashAttention優(yōu)化

模型在推理也就是解碼過程中，為了計(jì)算softmax(queries @keys.transpose)@values這兩個(gè)值，生成的每個(gè)新token都需要關(guān)注先前的所有token。

團(tuán)隊(duì)先前的工作FlashAttention，已經(jīng)在訓(xùn)練階段對(duì)此操作進(jìn)行了優(yōu)化。

當(dāng)時(shí)，F(xiàn)lashAttention解決的主要瓶頸是讀寫中間結(jié)果的內(nèi)存帶寬（例如，Q @ K^T）。

然而，在推理階段，我們要面對(duì)的瓶頸變了，導(dǎo)致FlashAttention所做的優(yōu)化并不能直接拿過來應(yīng)用。

具體而言：

在階段階段，F(xiàn)lashAttention在batch size和查詢長(zhǎng)度維度上進(jìn)行并行化。

在推理階段，查詢長(zhǎng)度通常為1，這意味著如果batch size小于GPU上的流式多處理器數(shù)量（例如，A100為108），該操作將僅使用GPU的一小部分。

這對(duì)于長(zhǎng)上下文情況尤甚，因?yàn)殚L(zhǎng)上下文需要較小的batch size才能適應(yīng)GPU內(nèi)存。

所以，結(jié)果就是，當(dāng)batch size為1時(shí)，F(xiàn)lashAttention將只占用不足1%的GPU，非常不劃算。

當(dāng)然，你可能會(huì)說，不用FlashAttention也行，用矩陣乘法原語來完注意力操作。

不過，作者指出，這種情況又會(huì)完全占用GPU，并啟動(dòng)非常多的寫入和讀取中間結(jié)果的內(nèi)核，也不是最佳辦法。

Flash-Decoding誕生

最終，基于以上考量，作者在FlashAttention的基礎(chǔ)上，添加了一個(gè)新的并行化緯度：key和value序列長(zhǎng)度。

這個(gè)方法（即Flash-Decoding）結(jié)合上述兩種方法的優(yōu)點(diǎn)：

與FlashAttention一樣，它在全局內(nèi)存中存儲(chǔ)的額外數(shù)據(jù)非常少，但只要上下文長(zhǎng)度足夠大，即使batch size很小，它也可以充分利用GPU。

詳細(xì)來看，F(xiàn)lash-Decoding一共分為三個(gè)步驟：

1、先將key和value值分成更小的塊。

2、用FlashAttention并行計(jì)算每塊分割的查詢注意力。并為每行和每塊分割寫入一個(gè)額外標(biāo)量：注意力值的log-sum-exp。

3、最后，通過減少所有分割來計(jì)算實(shí)際輸出，使用log-sum-exp來scale每塊分割的貢獻(xiàn)。

作者指出，由于attention/softmax可以迭代計(jì)算，以上所有操作均可行。

并且在Flash-Decoding中，ttention/softmax既可以在分割塊內(nèi)，也可以跨分割塊來執(zhí)行最終的縮減，只不過后者可縮減的步驟很少。

而在實(shí)際操作中，步驟1不涉及任何GPU操作，因?yàn)閗ey和value塊是完整的張量視圖。然后由2個(gè)獨(dú)立的內(nèi)核分別執(zhí)行步驟2和3。

最高提速8倍

驗(yàn)證環(huán)節(jié)，作者在CodeLLaMa-34b（架構(gòu)與Llama 2相同）上對(duì)其解碼吞吐量進(jìn)行了基準(zhǔn)測(cè)試。

具體以tok/s為單位，測(cè)量了512到64k序列長(zhǎng)度下的解碼速度（上限為從內(nèi)存中讀取整個(gè)模型以及KV緩存所需的時(shí)間），并和多種計(jì)算注意力的方法進(jìn)行對(duì)比，包括：

• Pytorch，使用純PyTorch原語運(yùn)行注意力
• FlashAttention v2
• FasterTransformer：使用FasterTransformer注意力內(nèi)核

最終，F(xiàn)lash-Decoding最高可將長(zhǎng)序列解碼速度提升8倍，并比其他方法具 有更好的擴(kuò)展性（受長(zhǎng)度影響較?。?/p>

此外，作者還在A100上對(duì)各種序列長(zhǎng)度和batch size的縮放多頭注意力進(jìn)行了微基準(zhǔn)測(cè)試。

結(jié)果顯示，當(dāng)序列長(zhǎng)度擴(kuò)展到64k時(shí)，Flash-Decoding實(shí)現(xiàn)了幾乎恒定的運(yùn)行時(shí)間。

如何使用？

以下是Flash-Decoding的獲取途徑，戳文末官方博客即可找到地址：

• FlashAttention包，2.2版本及以上
• xFormers包，0.0.22版本及以上

調(diào)度程序?qū)⒏鶕?jù)問題的大小自動(dòng)使用Flash-Decoding或 FlashAttention方法。

團(tuán)隊(duì)介紹

目前Flash-Decoding還沒出論文，但作者團(tuán)隊(duì)已透露，這次不再是Tri Dao“單打獨(dú)斗”，不過一作仍然是他。

Tri Dao今年博士畢業(yè)于斯坦福，7月份加盟大模型創(chuàng)業(yè)公司Together AI擔(dān)任首席科學(xué)家。

明年9月將上任普林斯頓大學(xué)助理教授，他是FlashAttention v1和v2的主要作者。

剩下三位作者分別是：

Daniel Haziza，F(xiàn)acebook AI Research研究工程師，主要負(fù)責(zé)xformers（用于訓(xùn)練加速的開源框架）；

Francisco Massa，同F(xiàn)acebook AI Research研究工程師， 主要從事PyTorch相關(guān)工作；

Grigory Sizov，Meta機(jī)器學(xué)習(xí)工程師，主要工作是優(yōu)化GPU上的LLM推理和其他AI工作負(fù)載，為PyTorch生態(tài)做出過貢獻(xiàn)。

官方博客：
https://princeton-nlp.github.io/flash-decoding/
參考鏈接：
https://twitter.com/tri_dao/status/1712904220519944411?s=20