微軟 RetrievalAttention: LLM+ANN, LLM 推理速度與精度的平衡
一、背景
本文我們繼續介紹一個針對超長上下文的 LLM 推理加速工作，同樣是 Token 稀疏化的方案，來解決 LLM 在超長序列場景計算量大、GPU 顯存消耗大的問題，不過結合了 ANN 檢索，可以實現更高的精度。
對應的論文為：[2409.10516] RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval
二、摘要
本文中作者提出了 RetrievalAttention，無需訓練就可以加速 Attention 計算。為了利用 Attention 的動態稀疏特性，RetrievalAttention 在 CPU 內存中使用 KV Cache 構建近似檢索（ANN）索引，并在生成過程中通過向量檢索識別最相關的索引。由于 Query 向量和 Key 向量之間存在 Out-Of-Distribution（OOD）問題，現成的 ANN 檢索仍然需要掃描 O(N) 數據（通常占所以 Key 的 30%）進行準確檢索，無法利用高稀疏性。
為了解決這個挑戰，RetrievalAttention 采用注意力感知向量檢索算法，可以調整 Query 只訪問 1-3% 的數據，從而實現亞線性時間復雜度。RetrievalAttention 大幅降低了長上下文 LLM 推理的成本，大幅降低 GPU 顯存需求，同時保持模型準確性。特別的，RetrievalAttention 只需要 16GB 內存就可以在具有 8B 參數的 LLM 上支持 128K Token的推理，在單個 RTX4090（24GB）上可以在 0.188s 內生成一個 Token。
如下圖 Figure 1 所示為本文方法與幾種常見方案的對比（PS：可以看出，本文方案相比之前 Token 稀疏化方案，是在犧牲一定推理速度的情況下提升精度）：
三、方法
3.1 背景
使用 ANN 來識別關鍵 Token 有個獨特的挑戰：當前大部分的 ANN 引擎都假設 Query 向量和 Key 向量滿足相同的分布，以此來實現高召回率。作者在這篇論文中首次提出這種假設在 Attention 機制中不成立。Query 的這種 OOD 特性損壞了 ANN 的預期檢索質量，從而導致不得不訪問更多的數據來保持正確性，作者實驗表明，為了維持可接受的準確率，至少需要掃描 30% 的 Key 向量。
如下圖 Figure 2 所示：
- （a）Attention Score 具有非常高的稀疏性，64000 個 Token，只有不到 500 Token 的 Score 大于 10-6。
- （b）Q 和 Q 或者 K 和 K 的相關性很高，而 Q 和 K 的相關性很差，需要掃描 30% 左右的 Token 才能保證 0.8 左右的召回率。
- （c）同樣說明了 Q 和 K 的距離比較遠。?
3.2 概覽
本文的工作主要聚焦于 Token Decoding 階段，會假設 Prefill 階段已經執行完成，比如通過 Context Caching 方案或 Prefill 和 Decoding 分離方案。
如下圖 Figure 3（a）所示為本文方案 RetrievalAttention 的概覽，其利用 CPU 側的 ANN 檢索來實現近似 Attention 計算，為了支持長序列，也會將所有 KV Cache Offload 到 CPU 內存以便構建索引。如圖（b）所示是為了解決 OOD 問題而采用的索引機制。
3.3 近似 Attention
具體來說，不使用完整的 Attention Score，而是采用最相關的 KV 向量來近似 Attention Score：
3.4 Attention 感知向量檢索
對于每對 Key 和 Value，首先確定是放在 CPU Memory 還是 GPU Memory（方法見下一小節）。然后 Offload 到 CPU 內存的 Key 和 Value 會使用 Key 來構建索引，并使用 Query 來檢索。
為了加速 Token 生成過程中的向量檢索速度，RetrievalAttention 利用 Prefill 階段的現有 Query 來指導 Key 向量的索引構建。如上圖 Figure 3（b）所示，RetrievalAttention 顯式的建立從 Query 向量到其最近的 Key 向量的連接（即精確的 K 個最近鄰，或 KNN）。KNN 結果可以通過 GPU 高效計算，形成從 Query 向量分布到 Key 向量分布的映射。使用這種結構，Decoding 的 Query 向量查詢時可以首先查詢最近的 Query 向量，然后將其映射為 Key 向量。
因此，之前的 Query 向量充當了解決 OOD 問題的橋梁。然而，這種結構在內存開銷和搜索效率方面仍然存在缺陷，因為除了 Key 向量之外，還需要存儲和訪問 Query 向量。為了解決這個問題，作者利用先進的跨模態 ANN 索引 RoarGraph 中的投影技術來消除 Query 向量。具體來說，通過使用 Query 向量和 Key 向量的連接關系，將 KNN 連接投影到 Key 向量中，從而有效地簡化搜索。此外，此方法也允許對未來的 Query 向量進行高效的索引遍歷。
作者實驗結果表明，通過這種 Query 和 Key 的連接關系進行有效建模，向量數據庫只需掃描 1-3% 的 Key 向量即可達到高召回率，與 IVF 索引相比，索引搜索延遲大幅降低 74%。
3.5 CPU 和 GPU 協同執行
為了利用 GPU 并行性加速注意力計算，RetrievalAttention 將注意力計算分解為兩組不相交的 KV Cache 向量：GPU 上的可預測向量和 CPU 上的動態向量，然后將兩部分 Attention 輸出合并在一起作為完整的 Attention 輸出。
具體來說，利用 Prefill 階段觀察到的模式來預測 Token 生成過程中持續激活的 KV 向量。與 StreamingLLM 類似，作者將固定的幾個初始 Token 和最近窗口內的 Token 作為靜態 Token，持久化在 GPU 上。RetrievalAttention 也可以適配更復雜的靜態模式，以便實現低推理成本和高準確性的平衡。為了最大限度減少通過慢速 PCIe 的數據傳輸，RetrievalAttention 在 CPU 和 GPU 上獨立計算 Attention，然后將其組合起來，這個靈感來自 FastAttention（[2205.14135] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness）。
四、評估
4.1 實驗配置
機器包含三個：
- RTX 4090 GPU（24G 顯存），Intel i9-10900X CPU（20 Core），128 GB 內存。
- A100 GPU（80GB 顯存），AMD EPYC CPU（24 Core）。
- A100 GPU（80GB 顯存），AMD EPYC 7V12 CPU（48 Core），1.72TB 內存。
模型包含三個：
- LLaMA-3-8B-Instruct-262K
- Yi-6B-200K
- Yi-9B-200K
對比框架包括：
- Full Attention 的 vLLM
- StreamingLLM
- SnapKV
- InfLLM
基準測試包括：
- ∞-Bench
- RULER
- Needle-in-a-haystack
4.2 長文本任務精度
如下圖 Table 2 所示，本文提出的 RetrievalAttention 明顯優于之前的方案，平均精度非常接近 Full Attention。當然，部分模型上 Flat（暴露檢索索引數據） 會略好于 RetrievalAttention，不過差距不大。
4.3 時延評估
如下圖 Table 6 所示，作者首先驗證了本文提出的檢索方式的有效性，可以看出，提出的 RetrievalAttention 相比 Flat 和 IVF 可以提供 4.9x 和 1.98x 的加速，證明了檢索機制的有效性：
如下圖 Table 4 所示，作者也與之前的其他稀疏化方案進行對比，可以看出，之前的方案往往采用固定的 Token 數，因此隨著序列變長并沒有明顯增加時延，而本文的方法會略微增加。同時，本文方法推理 Latency 相比之前方法明顯增加，大概是之前方法 Latency 的 3x-6x。然而其仍然明顯低于 Full Attention（FlexGen） 的結果。相當于在效果和速度之間的折衷。
如下圖 Table 7 和 Table 8 為在 A100 上的結果，結論類似，不過在 100K 和 200K 時其 Latency 會超過 vLLM：
五、參考鏈接
本文轉載自 ??AI閑談??，作者： AI閑談
