今天小編分享的科學經驗：把注意力計算丢給CPU，大模型解碼吞吐量提高1.76~4.99倍，歡迎閱讀。
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 Attention 的問題
眾所周知，注意力機制本質上具有稀疏性，因此動态稀疏注意力和基于 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）允許從一個已知分布中抽取樣本 1，2，…，B，來估計未知分布的期望。
最終的輸出由下式給出：
重要性采樣要求和的峰值對應以降低估計方差，為此，研究使用局部敏感哈希（LSH） 來生成采樣概率。
需要指出的是，因為存在 Softmax（注意力得分需要歸一化）, 所以研究實際上試圖近似的是自歸一化重要性采樣。
系統：将注意力計算和哈希表放在 CPU 上
除了精度下降的問題外，受限的 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 哈希表，而無需擔心哈希碼計算和哈希表存儲的開銷。
實驗
研究從準确率和推理速度兩個方面來評估 MagicPIG 系統的能力。
圖片中的百分比為實際采樣的 KV cache 的數量，對于 MagicPIG 而言，K10L150 ≈ 2%, K10L170 ≈ 2.5%。
長文本 RULER
以 Llama-3.1-8B-Instruct 為例，MagicPIG 在檢索和推理任務中比 Quest（稀疏注意力的 SOTA 基線）實現了更高的下遊準确率。
推理速度和吞吐量
在 L20 + Intel 8563C 上測試吞吐量，MagicPIG 與僅使用 GPU 的注意力機制相比，在各種情況下提高了 1.76~4.99 倍的解碼吞吐量。
整體而言，MagicPIG 是将經典的哈希算法和高維向量估計用到 LLM 解碼上的嘗試。
接下來，研究将支持更加高效的局部敏感哈希算法，并希望進一步降低 LLM 部署成本，探索異構計算的可能性。
論文：
https://arxiv.org/abs/2410.16179
項目地址：
www.lsh-ai.com
— 完 —
投稿請發郵件到：
标題注明【投稿】，告訴我們：
你是誰，從哪來，投稿内容
附上論文 / 項目主頁鏈接，以及聯系方式哦
我們會（盡量）及時回復你
點這裡關注我，記得标星哦～
一鍵三連「分享」、「點贊」和「在看」
科技前沿進展日日相見 ~
>
