這兩天，FlashAttention團隊推出了新作：

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

該方法尤其造福于長上下文LLM，在64k長度的CodeLlama-34B上通過了驗證。

甚至得到了PyTorch官方認可：

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

不過它僅限于訓練階段。

因此，這一新成果一出，就有網友表示：

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

據介紹，這個新方法也是在FlashAttention的基礎之上衍生而出，主要思想也不復雜：

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

提速8倍的長上下文推理方法來了

該方法被命名為Flash-Decoding。

背景與動機

根據作者介紹：

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

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

但有一個操作例外：注意力 (attention)，它不能隨著上下文長度靈活擴展。

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

更別提attention在batch size上進行擴展時，即使模型上下文相對較短，它也可能成為性能瓶頸（因為模型要讀取的內存量與batch size成比例，而它僅取決于模型其余部分的大小）。

怎么解除？

不可復用的FlashAttention優化

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

團隊先前的工作FlashAttention，已經在訓練階段對此操作進行了優化。

當時，FlashAttention解決的主要瓶頸是讀寫中間結果的內存帶寬（例如，Q @ K^T）。

然而，在推理階段，我們要面對的瓶頸變了，導致FlashAttention所做的優化并不能直接拿過來應用。

具體而言：

在階段階段，FlashAttention在batch size和查詢長度維度上進行并行化。

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

這對于長上下文情況尤甚，因為長上下文需要較小的batch size才能適應GPU內存。

所以，結果就是，當batch size為1時，FlashAttention將只占用不足1%的GPU，非常不劃算。

當然，你可能會說，不用FlashAttention也行，用矩陣乘法原語來完注意力操作。

不過，作者指出，這種情況又會完全占用GPU，并啟動非常多的寫入和讀取中間結果的內核，也不是最佳辦法。

Flash-Decoding誕生

最終，基于以上考量，作者在FlashAttention的基礎上，添加了一個新的并行化緯度：key和value序列長度。

這個方法（即Flash-Decoding）結合上述兩種方法的優點：

與FlashAttention一樣，它在全局內存中存儲的額外數據非常少，但只要上下文長度足夠大，即使batch size很小，它也可以充分利用GPU。

詳細來看，Flash-Decoding一共分為三個步驟：

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

2、用FlashAttention并行計算每塊分割的查詢注意力。并為每行和每塊分割寫入一個額外標量：注意力值的log-sum-exp。

3、最后，通過減少所有分割來計算實際輸出，使用log-sum-exp來scale每塊分割的貢獻。

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

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

而在實際操作中，步驟1不涉及任何GPU操作，因為key和value塊是完整的張量視圖。然后由2個獨立的內核分別執行步驟2和3。

最高提速8倍

驗證環節，作者在CodeLLaMa-34b（架構與Llama 2相同）上對其解碼吞吐量進行了基準測試。

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

• Pytorch，使用純PyTorch原語運行注意力
• FlashAttention v2
• FasterTransformer：使用FasterTransformer注意力內核

     最終，Flash-Decoding最高可將長序列解碼速度提升8倍，并比其他方法具       有更好的擴展性（受長度影響較小）

此外，作者還在A100上對各種序列長度和batch size的縮放多頭注意力進行了微基準測試。

結果顯示，當序列長度擴展到64k時，Flash-Decoding實現了幾乎恒定的運行時間。

如何使用？

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

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

調度程序將根據問題的大小自動使用Flash-Decoding或 FlashAttention方法。

團隊介紹

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

Tri Dao今年博士畢業于斯坦福，7月份加盟大模型創業公司Together AI擔任首席科學家。

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

剩下三位作者分別是：

Daniel Haziza，Facebook AI Research研究工程師，主要負責xformers（用于訓練加速的開源框架）；

Francisco Massa，同Facebook AI Research研究工程師， 主要從事PyTorch相關工作；

Grigory Sizov，Meta機器學習工程師，主要工作是優化GPU上的LLM推理和其他AI工作負載，為PyTorch生態做出過貢獻。

官方博客：https://princeton-nlp.github.io/flash-decoding/