別再「浪費」GPU了，FlashAttention重磅升級，實現長文本推理速度8倍提升
最近，像 ChatGPT 或 Llama 這樣的大型語言模型（LLM）引起了前所未有的關注。然而，它們的運行成本仍然極高。雖然生成單個響應可能僅需 0.01 美元（在 AWS 上的 8xA100 實例上運行幾秒鐘），但當擴大規模以滿足數十億用戶的需求時，成本會迅速累積。而且，這些用戶可能每天與 LLM 進行多次互動。某些用例的成本更高，例如代碼自動生成，因為它會隨著每次輸入新字符而運行。隨著 LLM 應用的不斷增加，即使在生成時間方面實現細微的效率提升，也將產生巨大的影響。
LLM 推理（或「解碼」）是一個迭代的過程：token 逐個生成。生成包含 N 個 token 的完整句子需要通過模型進行 N 次前向傳遞。幸運的是，我們可以緩存先前計算的 token：這意味著單個生成步驟不依賴于上下文長度，除了一個單獨的操作 —— 注意力。這個操作導致上下文長度不能很好地擴展。
在 LLM 的重要新興用例中，有一些需要利用更長的上下文。只有擁有了更長的上下文窗口，LLM 才能對更長的文檔進行推理，無論是總結文檔還是回答其中的問題。此外，它們還可以保持更長的對話歷史，甚至在編寫代碼之前處理整個代碼庫。舉個例子，在 2022 年，大多數 LLM 的上下文長度最多為 2k（例如 GPT-3），但現在，有些開源 LLM 已經可以擴展到 32k（比如 Llama-2-32k），甚至有些模型已經達到了 100k（比如 CodeLlama）。在這些情境中，注意力操作在推理過程中占據了相當大的時間比例。
在擴展 batch size 維度時，即使上下文相對較短，注意力也可能成為一個瓶頸。這是因為隨著 batch 維度的增加，需要讀取的內存量也會增加，而對于模型的其余部分，內存需求只取決于模型的大小。
為了解決上述問題，FlashAttention 的作者 Tri Dao 等人提出了一項名為「Flash-Decoding」的技術，它顯著加速了推理過程中的注意力計算，使長序列的處理生成速度提高到了原來的 8 倍。其主要思想是以最快的速度并行加載鍵和值，然后分別重新縮放和合并結果，以維持正確的注意力輸出。
解碼時的多頭注意力
在解碼期間，生成的每個新 token 都需要關注所有先前的 token，以計算：softmax (queries @ keys.transpose) @ values
這個操作已經在訓練階段通過 FlashAttention 進行了優化（包括最近的 v1 和 v2 版本），瓶頸是讀寫中間結果的內存帶寬（如 Q @ K^T）。然而，這些優化并不直接適用于推理情況，因為瓶頸不同。在訓練中，FlashAttention 并行處理 batch size 和查詢長度兩個維度。而在推理過程中，查詢長度通常為 1：這意味著，如果 batch size 小于 GPU 上的流多處理器（streaming multiprocessor，SM）數量（例如 A100 有 108 個），該操作只會利用 GPU 的一小部分！特別是在處理長上下文時，情況尤為明顯，因為它需要較小的 batch size 以適應 GPU 內存。當 batch size 為 1 時，FlashAttention 將使用不到 1% 的 GPU！
FlashAttention 只在查詢塊和 batch size 之間并行，并且在解碼期間不會設法占用整個 GPU
使用矩陣乘法基元也能執行注意力計算，這樣就不需要使用 FlashAttention 了。在這種情況下，該操作會占用整個 GPU，但會啟動許多寫入和讀取中間結果的內核，因此并不是最優的做法。
更快的注意力解碼：Flash-Decoding
新方法 Flash-Decoding 基于 FlashAttention，同時引入了一個新的并行維度：鍵值序列的長度。它綜合了上述兩種方法的優點。與 FlashAttention 類似，它在全局內存中存儲的額外數據很少。然而，只要上下文足夠長，即使 batch size 較小，它也能充分利用 GPU。
Flash-Decoding 也在鍵和值之間并行化，代價是一個小的最終歸約（reduction 步驟。
Flash-Decoding 主要有三個工作步驟：
- 首先，將鍵 / 值分成更小的塊；
- 使用 FlashAttention 并行計算查詢與每個這些分塊的注意力，為每行和每個分塊額外寫入一個標量值：注意力值的 log-sum-exp
- 最后，通過對所有分塊進行歸約來計算實際輸出，使用 log-sum-exp 來調整每個分塊的貢獻。
這一切之所以可行，都是因為注意力 /softmax 可以進行迭代計算。在 Flash-Decoding 中，它在兩個級別上被使用：在分塊內部（類似 FlashAttention），以及跨分塊進行最終的歸約計算。
實際操作中，步驟（1）不涉及任何 GPU 操作，因為鍵 / 值塊是完整鍵 / 值張量的視圖。然后，有兩個獨立的核函數，分別用于執行步驟（2）和（3）。
在 CodeLlama 34B 上進行的基準測試
為了驗證上述新方法，研究者對 CodeLLaMa-34b 的解碼吞吐量進行了基準測試。該模型與 Llama 2 具有相同的架構，一般來說，結果應該適用于許多大型語言模型。研究者在不同序列長度下（從 512 到 64k），以 tok/s 為單位來測量解碼速度，并比較了多種計算注意力的方式：
- Pytorch：使用純粹的 PyTorch 基元來運行注意力計算（不使用 FlashAttention）；
- FlashAttention v2；
- FasterTransformer：使用 FasterTransformer 的注意力內核；
- Flash-Decoding；
以及一個上限值，該值計算了從內存中讀取整個模型和 KV-cache 所需的時間
對于非常大的序列，Flash-Decoding 可以將解碼速度提高至 8 倍，并且比其他方法的擴展性要好得多。
在 prompt 比較小時，所有方法表現接近。但是當序列長度從 512 增加到 64k 時，除了 Flash-Decoding，其他方法的可擴展性都很差。在 Flash-Decoding 的這種模式下（batch size 為 1），擴展序列長度對生成速度的影響很小。
組件級微基準測試
研究者還在 A100 上對多頭注意力進行了微基準測試，輸入為 f16，考慮了不同的序列長度和 batch size。他們將 batch size 設置為 1，并且使用 16 個 128 維的查詢頭，以及 2 個鍵 / 值頭（分組查詢注意力），這與在 4 個 GPU 上運行的 CodeLLaMa-34b 使用的維度相匹配。
上述微基準測試展示了多頭注意力的運行時間，單位為微秒。Flash-Decoding 在序列長度擴展到高達 64k 時，幾乎實現了恒定的運行時間。
之前測量的高達 8 倍的端到端加速是可能的，因為注意力本身的速度比 FlashAttention 快高達 50 倍。在序列長度達到 32k 之前，注意力的時間大致是恒定的，因為 Flash-Decoding 能夠完全利用 GPU。
使用 Flash-Decoding
Flash-decoding 可以在以下鏈接中找到：
- FlashAttention 包，從 v2.2 開始：https://github.com/Dao-AILab/flash-attention/tree/main
- xFormers 包（搜索 xformers.ops.memory_efficient_attention），從 0.0.22 開始：調度程序將根據問題的大小自動使用 Flash-Decoding 或 FlashAttention 方法。當這些方法不受支持時，它可以調度到一個高效的 triton 內核，該內核實現了 Flash-Decoding 算法。
一個完整的使用 LLaMa v2 / CodeLLaMa 的解碼示例可以在 FlashAttention repo 和 xFormers repo 中找到。此外，作者還提供了一個簡單的 LLaMa v1/v2 模型的高效解碼代碼示例，旨在快速、易讀、有教育意義和易于修改。
參考鏈接：https://princeton-nlp.github.io/flash-decoding/
