最近，像 ChatGPT 或 Llama 這樣的大型語(yǔ)言模型（LLM）引起了前所未有的關(guān)注。然而，它們的運(yùn)行成本仍然極高。雖然生成單個(gè)響應(yīng)可能僅需 0.01 美元（在 AWS 上的 8xA100 實(shí)例上運(yùn)行幾秒鐘），但當(dāng)擴(kuò)大規(guī)模以滿足數(shù)十億用戶的需求時(shí)，成本會(huì)迅速累積。而且，這些用戶可能每天與 LLM 進(jìn)行多次互動(dòng)。某些用例的成本更高，例如代碼自動(dòng)生成，因?yàn)樗鼤?huì)隨著每次輸入新字符而運(yùn)行。隨著 LLM 應(yīng)用的不斷增加，即使在生成時(shí)間方面實(shí)現(xiàn)細(xì)微的效率提升，也將產(chǎn)生巨大的影響。

LLM 推理（或「解碼」）是一個(gè)迭代的過(guò)程：token 逐個(gè)生成。生成包含 N 個(gè) token 的完整句子需要通過(guò)模型進(jìn)行 N 次前向傳遞。幸運(yùn)的是，我們可以緩存先前計(jì)算的 token：這意味著單個(gè)生成步驟不依賴于上下文長(zhǎng)度，除了一個(gè)單獨(dú)的操作 —— 注意力。這個(gè)操作導(dǎo)致上下文長(zhǎng)度不能很好地?cái)U(kuò)展。

在 LLM 的重要新興用例中，有一些需要利用更長(zhǎng)的上下文。只有擁有了更長(zhǎng)的上下文窗口，LLM 才能對(duì)更長(zhǎng)的文檔進(jìn)行推理，無(wú)論是總結(jié)文檔還是回答其中的問(wèn)題。此外，它們還可以保持更長(zhǎng)的對(duì)話歷史，甚至在編寫(xiě)代碼之前處理整個(gè)代碼庫(kù)。舉個(gè)例子，在 2022 年，大多數(shù) LLM 的上下文長(zhǎng)度最多為 2k（例如 GPT-3），但現(xiàn)在，有些開(kāi)源 LLM 已經(jīng)可以擴(kuò)展到 32k（比如 Llama-2-32k），甚至有些模型已經(jīng)達(dá)到了 100k（比如 CodeLlama）。在這些情境中，注意力操作在推理過(guò)程中占據(jù)了相當(dāng)大的時(shí)間比例。

在擴(kuò)展 batch size 維度時(shí)，即使上下文相對(duì)較短，注意力也可能成為一個(gè)瓶頸。這是因?yàn)殡S著 batch 維度的增加，需要讀取的內(nèi)存量也會(huì)增加，而對(duì)于模型的其余部分，內(nèi)存需求只取決于模型的大小。

為了解決上述問(wèn)題，F(xiàn)lashAttention 的作者 Tri Dao 等人提出了一項(xiàng)名為「Flash-Decoding」的技術(shù)，它顯著加速了推理過(guò)程中的注意力計(jì)算，使長(zhǎng)序列的處理生成速度提高到了原來(lái)的 8 倍。其主要思想是以最快的速度并行加載鍵和值，然后分別重新縮放和合并結(jié)果，以維持正確的注意力輸出。

解碼時(shí)的多頭注意力

在解碼期間，生成的每個(gè)新 token 都需要關(guān)注所有先前的 token，以計(jì)算：softmax (queries @ keys.transpose) @ values

這個(gè)操作已經(jīng)在訓(xùn)練階段通過(guò) FlashAttention 進(jìn)行了優(yōu)化（包括最近的 v1 和 v2 版本），瓶頸是讀寫(xiě)中間結(jié)果的內(nèi)存帶寬（如 Q @ K^T）。然而，這些優(yōu)化并不直接適用于推理情況，因?yàn)槠款i不同。在訓(xùn)練中，F(xiàn)lashAttention 并行處理 batch size 和查詢長(zhǎng)度兩個(gè)維度。而在推理過(guò)程中，查詢長(zhǎng)度通常為 1：這意味著，如果 batch size 小于 GPU 上的流多處理器（streaming multiprocessor，SM）數(shù)量（例如 A100 有 108 個(gè)），該操作只會(huì)利用 GPU 的一小部分！特別是在處理長(zhǎng)上下文時(shí)，情況尤為明顯，因?yàn)樗枰^小的 batch size 以適應(yīng) GPU 內(nèi)存。當(dāng) batch size 為 1 時(shí)，F(xiàn)lashAttention 將使用不到 1% 的 GPU！

FlashAttention 只在查詢塊和 batch size 之間并行，并且在解碼期間不會(huì)設(shè)法占用整個(gè) GPU

使用矩陣乘法基元也能執(zhí)行注意力計(jì)算，這樣就不需要使用 FlashAttention 了。在這種情況下，該操作會(huì)占用整個(gè) GPU，但會(huì)啟動(dòng)許多寫(xiě)入和讀取中間結(jié)果的內(nèi)核，因此并不是最優(yōu)的做法。

更快的注意力解碼：Flash-Decoding

新方法 Flash-Decoding 基于 FlashAttention，同時(shí)引入了一個(gè)新的并行維度：鍵值序列的長(zhǎng)度。它綜合了上述兩種方法的優(yōu)點(diǎn)。與 FlashAttention 類似，它在全局內(nèi)存中存儲(chǔ)的額外數(shù)據(jù)很少。然而，只要上下文足夠長(zhǎng)，即使 batch size 較小，它也能充分利用 GPU。

Flash-Decoding 也在鍵和值之間并行化，代價(jià)是一個(gè)小的最終歸約（reduction 步驟。

Flash-Decoding 主要有三個(gè)工作步驟：

• 首先，將鍵 / 值分成更小的塊；
• 使用 FlashAttention 并行計(jì)算查詢與每個(gè)這些分塊的注意力，為每行和每個(gè)分塊額外寫(xiě)入一個(gè)標(biāo)量值：注意力值的 log-sum-exp
• 最后，通過(guò)對(duì)所有分塊進(jìn)行歸約來(lái)計(jì)算實(shí)際輸出，使用 log-sum-exp 來(lái)調(diào)整每個(gè)分塊的貢獻(xiàn)。

這一切之所以可行，都是因?yàn)樽⒁饬?/softmax 可以進(jìn)行迭代計(jì)算。在 Flash-Decoding 中，它在兩個(gè)級(jí)別上被使用：在分塊內(nèi)部（類似 FlashAttention），以及跨分塊進(jìn)行最終的歸約計(jì)算。

實(shí)際操作中，步驟（1）不涉及任何 GPU 操作，因?yàn)殒I / 值塊是完整鍵 / 值張量的視圖。然后，有兩個(gè)獨(dú)立的核函數(shù)，分別用于執(zhí)行步驟（2）和（3）。

在 CodeLlama 34B 上進(jìn)行的基準(zhǔn)測(cè)試

為了驗(yàn)證上述新方法，研究者對(duì) CodeLLaMa-34b 的解碼吞吐量進(jìn)行了基準(zhǔn)測(cè)試。該模型與 Llama 2 具有相同的架構(gòu)，一般來(lái)說(shuō)，結(jié)果應(yīng)該適用于許多大型語(yǔ)言模型。研究者在不同序列長(zhǎng)度下（從 512 到 64k），以 tok/s 為單位來(lái)測(cè)量解碼速度，并比較了多種計(jì)算注意力的方式：

• Pytorch：使用純粹的 PyTorch 基元來(lái)運(yùn)行注意力計(jì)算（不使用 FlashAttention）；
• FlashAttention v2；
• FasterTransformer：使用 FasterTransformer 的注意力內(nèi)核；
• Flash-Decoding；

以及一個(gè)上限值，該值計(jì)算了從內(nèi)存中讀取整個(gè)模型和 KV-cache 所需的時(shí)間

對(duì)于非常大的序列，F(xiàn)lash-Decoding 可以將解碼速度提高至 8 倍，并且比其他方法的擴(kuò)展性要好得多。

在 prompt 比較小時(shí)，所有方法表現(xiàn)接近。但是當(dāng)序列長(zhǎng)度從 512 增加到 64k 時(shí)，除了 Flash-Decoding，其他方法的可擴(kuò)展性都很差。在 Flash-Decoding 的這種模式下（batch size 為 1），擴(kuò)展序列長(zhǎng)度對(duì)生成速度的影響很小。

組件級(jí)微基準(zhǔn)測(cè)試

研究者還在 A100 上對(duì)多頭注意力進(jìn)行了微基準(zhǔn)測(cè)試，輸入為 f16，考慮了不同的序列長(zhǎng)度和 batch size。他們將 batch size 設(shè)置為 1，并且使用 16 個(gè) 128 維的查詢頭，以及 2 個(gè)鍵 / 值頭（分組查詢注意力），這與在 4 個(gè) GPU 上運(yùn)行的 CodeLLaMa-34b 使用的維度相匹配。

上述微基準(zhǔn)測(cè)試展示了多頭注意力的運(yùn)行時(shí)間，單位為微秒。Flash-Decoding 在序列長(zhǎng)度擴(kuò)展到高達(dá) 64k 時(shí)，幾乎實(shí)現(xiàn)了恒定的運(yùn)行時(shí)間。

之前測(cè)量的高達(dá) 8 倍的端到端加速是可能的，因?yàn)樽⒁饬Ρ旧淼乃俣缺?FlashAttention 快高達(dá) 50 倍。在序列長(zhǎng)度達(dá)到 32k 之前，注意力的時(shí)間大致是恒定的，因?yàn)?Flash-Decoding 能夠完全利用 GPU。

使用 Flash-Decoding

Flash-decoding 可以在以下鏈接中找到：

• FlashAttention 包，從 v2.2 開(kāi)始：https://github.com/Dao-AILab/flash-attention/tree/main
• xFormers 包（搜索 xformers.ops.memory_efficient_attention），從 0.0.22 開(kāi)始：調(diào)度程序?qū)⒏鶕?jù)問(wèn)題的大小自動(dòng)使用 Flash-Decoding 或 FlashAttention 方法。當(dāng)這些方法不受支持時(shí)，它可以調(diào)度到一個(gè)高效的 triton 內(nèi)核，該內(nèi)核實(shí)現(xiàn)了 Flash-Decoding 算法。

一個(gè)完整的使用 LLaMa v2 / CodeLLaMa 的解碼示例可以在 FlashAttention  repo 和 xFormers  repo 中找到。此外，作者還提供了一個(gè)簡(jiǎn)單的 LLaMa v1/v2 模型的高效解碼代碼示例，旨在快速、易讀、有教育意義和易于修改。

參考鏈接：https://princeton-nlp.github.io/flash-decoding/