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

基于局部敏感哈希采樣技術(shù)

原標(biāo)題：把注意力計(jì)算丟給CPU，大模型解碼吞吐量提高1.76~4.99倍
文章來源：量子位
內(nèi)容字?jǐn)?shù)：5268字

MagicPIG: 緩解大語言模型KV緩存壓力的創(chuàng)新方案

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

1. KV緩存瓶頸與現(xiàn)有方法的不足

1. 在長(zhǎng)上下文LLM推理中，KV緩存用于存儲(chǔ)中間注意力鍵和值，避免重復(fù)計(jì)算，但其內(nèi)存占用隨批量大小和序列長(zhǎng)度線性增長(zhǎng)，嚴(yán)重限制了GPU的批量處理能力。

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

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

2. MagicPIG的創(chuàng)新之處

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

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

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

3. 系統(tǒng)設(shè)計(jì)與實(shí)驗(yàn)結(jié)果

1. MagicPIG將LLM解碼分為四個(gè)部分：參數(shù)計(jì)算(GPU)、注意力計(jì)算(CPU)、隨機(jī)投影(GPU)和檢索(CPU)。這種任務(wù)分區(qū)有效地利用了CPU和GPU的優(yōu)勢(shì)。

2. 實(shí)驗(yàn)結(jié)果表明，MagicPIG在Llama-3.1-8B-Instruct模型上，相比于Quest，在檢索和推理任務(wù)中實(shí)現(xiàn)了更高的下游準(zhǔn)確率。同時(shí)，MagicPIG的解碼吞吐量提高了1.76~4.99倍。

4. 總結(jié)

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

聯(lián)系作者

文章來源：量子位
作者微信：
作者簡(jiǎn)介：追蹤人工智能新趨勢(shì)，關(guān)注科技行業(yè)新突破