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

基于局部敏感哈希采樣技術

原標題：把注意力計算丟給CPU，大模型解碼吞吐量提高1.76~4.99倍
文章來源：量子位
內容字數：5268字

MagicPIG: 緩解大語言模型KV緩存壓力的創新方案

來自卡內基梅隆大學、華盛頓大學和Meta AI的研究人員提出了一種名為MagicPIG的新方法，有效解決了大語言模型(LLM)推理過程中GPU內存容量限制的問題，該問題主要由KV緩存的巨大內存占用引起。

1. KV緩存瓶頸與現有方法的不足

1. 在長上下文LLM推理中，KV緩存用于存儲中間注意力鍵和值，避免重復計算，但其內存占用隨批量大小和序列長度線性增長，嚴重限制了GPU的批量處理能力。

2. 現有技術如TopK Attention、Quest、H2O和Loki等，試圖通過選擇注意力得分最高的子集來壓縮KV緩存，但這些方法存在精度損失的問題，尤其在需要高上下文利用率的復雜任務中，性能下降嚴重。

3. 研究人員觀察到注意力機制的三個關鍵特性，解釋了TopK方法的局限性：首個輸入token的隱藏狀態幾乎不隨輸入變化；鍵狀態的中心方向在不同輸入句子中保持穩定；鍵狀態的中心與匯聚點token的鍵狀態幾乎相反。

2. MagicPIG的創新之處

1. **基于采樣而非搜索:** MagicPIG采用基于采樣的注意力估計方法，而非像TopK那樣搜索最高的鍵值對。這通過引入基礎分布信息，顯著提高了估計精度，并降低了估計誤差。

2. **局部敏感哈希(LSH)技術:** MagicPIG利用LSH生成采樣概率，并結合重要性采樣技術，有效地從已知分布中抽取樣本，來估計未知分布的期望，從而降低估計方差。

3. **異構計算:** MagicPIG將注意力計算和哈希表卸載到CPU上，充分利用CPU的內存資源，降低了對GPU顯存的依賴。實驗表明，這種方法等效地提升了CPU的內存帶寬，在維持精度的情況下實現了高效的注意力計算。

3. 系統設計與實驗結果

1. MagicPIG將LLM解碼分為四個部分：參數計算(GPU)、注意力計算(CPU)、隨機投影(GPU)和檢索(CPU)。這種任務分區有效地利用了CPU和GPU的優勢。

2. 實驗結果表明，MagicPIG在Llama-3.1-8B-Instruct模型上，相比于Quest，在檢索和推理任務中實現了更高的下游準確率。同時，MagicPIG的解碼吞吐量提高了1.76~4.99倍。

4. 總結

MagicPIG通過巧妙地結合LSH采樣技術和異構計算，有效緩解了LLM推理過程中的KV緩存壓力，在保證精度的同時大幅提升了推理速度和吞吐量。這項研究為降低LLM部署成本，探索異構計算的可能性提供了新的思路。

聯系作者

文章來源：量子位
作者微信：
作者簡介：追蹤人工智能新趨勢，關注科技行業新突破