把注意力計算丟給CPU，大模型解碼吞吐量提高1.76~4.99倍

作者：量子位 2024-12-27 13:59:33

MagicPIG是將經典的哈希算法和高維向量估計用到LLM解碼上的嘗試。

CPU+GPU，模型KV緩存壓力被緩解了。

來自CMU、華盛頓大學、Meta AI的研究人員提出MagicPIG，通過在CPU上使用LSH（局部敏感哈希）采樣技術，有效克服了GPU內存容量限制的問題。

與僅使用GPU的注意力機制相比，MagicPIG在各種情況下提高了1.76~4.99倍的解碼吞吐量，并在檢索和推理任務中實現了更高的下游準確率，優于Quest等現有技術。

概括而言，這項研究主要貢獻有兩點：

1、相比于其他的稀疏注意力（Sparse Attention），MagicPIG基于采樣/估計而非搜索，提升了推理質量。

2、研究把解碼階段注意力模塊的計算和哈希表卸載到CPU上，探索了異構計算的可能性，并且提升了吞吐量，有望降低實際模型部署成本。

下面具體來看。

KV緩存限制了GPU高效利用

在長上下文大模型（LLM）的推理過程中，KV緩存（Key-Value Cache）成為關鍵瓶頸。KV緩存主要用于存儲中間的注意力鍵和值，從而避免重復計算。

然而，其顯存占用隨著批量大小和序列長度的線性增長而迅速增加，這嚴重限制了GPU的批量處理能力，導致計算資源無法被充分利用。

以NVIDIA A100-40GB GPU為例，在處理Llama-3.1-8B模型且上下文長度為128k時，僅支持單個請求，且近一半的解碼時間都消耗在訪問KV緩存上，GPU利用率明顯不足。

此外，推理過程中采用的一些策略，如多樣性生成（Best-of-N）和長鏈式推理（Long Chain-of-Thoughts），會進一步增加生成的Token數量，加劇顯存壓力，導致推理效率進一步下降。

眾所周知，注意力機制本質上具有稀疏性，因此動態稀疏注意力和基于TopK的近似方法得到了廣泛研究。

然而，這些方法往往伴隨著顯著的質量下降問題。

目前已有的KV緩存壓縮技術，如Quest、H2O和Loki，主要通過篩選出KV緩存中注意力得分最高的子集來提高效率。然而，盡管這些方法在實踐中表現出一定的效果，基于TopK的注意力依然是一種存在偏差的近似方法，且缺乏理論上的嚴格保障。

這種不足限制了其在高精度場景中的廣泛應用。

下圖顯示，即使是精確的TopK注意力機制也會導致顯著的估計誤差和下游任務性能下降。

這一問題在需要高上下文利用率的復雜任務中尤為突出，例如聚合任務、常用詞提取（CWE）、高頻詞提取（FWE）以及邏輯推理任務。在這些場景中，基于TopK近似方法的性能下降尤其嚴重。

以下幾點觀察揭示了為何TopK注意力機制無法始終有效工作。

這些觀察不僅解釋了注意力機制的行為，還可能對模型訓練具有重要意義：

1、首個輸入token（注意力匯聚點，sink）的隱藏狀態（包括但不限于鍵和值狀態）幾乎不隨輸入變化而改變。（見左圖，在采樣的輸入中，其最小相似度均高于0.99）

2、鍵狀態的中心方向在不同輸入句子中保持穩定。（見中圖，相似度均高于0.9）

3、鍵狀態的中心與匯聚點token的鍵狀態幾乎相反。（見右圖， -0.9至-0.8之間）

這些現象為理解注意力機制提供了新的視角，同時也表明傳統的TopK近似方法在某些場景下可能存在局限性。

為了解決這一問題，研究提出了一種基于采樣而非搜索TopK鍵值緩存的新方法。

與僅依賴注意力分數最高的鍵值對相比，融入基礎分布信息可以顯著提高估計的準確性。

研究將這一問題視為采樣中的偏差校正問題。在生物學、社會學和機器學習等領域，無偏且高效的采樣技術已被廣泛研究，并具有堅實的理論保障。

如圖所示，基于注意力分數按比例進行采樣（即所謂的Oracle Sampling，研究把注意力模塊的輸出看成value向量的期望值，對應的分布是注意力得分）相比于傳統的TopK選擇方法，其估計誤差要小得多，最多可降低4倍。

這表明采樣技術在注意力近似中的潛力。

從注意力得分??中采樣，在實際中不可行。重要性采樣（Importance Sampling）允許從一個已知分布??中抽取樣本??₁，??₂，…，??_B，來估計未知分布??的期望。

最終的輸出由下式給出：

重要性采樣要求??和??的峰值對應以降低估計方差，為此，研究使用局部敏感哈希（LSH）來生成采樣概率??。

需要指出的是，因為存在Softmax（注意力得分需要歸一化）, 所以研究實際上試圖近似的是自歸一化重要性采樣。

除了精度下降的問題外，受限的GPU顯存容量也限制了現有動態KV緩存壓縮方法（如Quest和Loki）在許多場景中的適用性。

與此同時，像DeepSpeed-Zero-Inference和FastDecode這樣的技術展示了將KV緩存和注意力計算卸載到CPU上的潛力。

CPU的內存帶寬大約是GPU顯存帶寬的10%-20%，這引出了一個自然的問題：

能否在不犧牲精度的前提下，將注意力計算中的內存訪問量減少10倍？

通過利用采樣算法，例如MagicPIG中基于LSH（局部敏感哈希）的采樣技術進行注意力估計，研究大幅降低了內存訪問量。這種方法等效地提升了CPU的內存帶寬，使得在維持精度的情況下實現高效的注意力計算。

論文的系統設計擴展了以往的工作，將大語言模型（LLM）的解碼分為以下四個部分：

參數計算：包括所有線性投均在GPU上運行。
注意力計算：涉及公式，該部分在CPU上運行。
隨機投影：在生成過程中，對于每個??執行K x L次隨機投影以生成哈希碼。由于所有注意力頭可以共享相同的隨機投影器，內存開銷較小（在實際實現中約為400KB）。實驗中K=9或10，而L為數百，因此該步驟主要受計算限制，放置在GPU上運行。
檢索：需要在L個哈希表中查找q的哈希碼。這部分計算開銷非常輕量，但預構建的哈希表占用的內存較大，因此更適合放置在CPU上運行。通過上述任務分區，可以支持更大規模的K和L哈希表，而無需擔心哈希碼計算和哈希表存儲的開銷。