Cache Me If You Can：陳丹琦團隊如何「抓住」關鍵緩存，解放LLM內存？

2025-06-25 09:15:38

陳丹琦團隊提出了「KV 足跡」作為一種統一的度量標準，它是在所有時間步中，未被逐出的鍵值緩存條目所占比例的聚合值。

普林斯頓大學計算機科學系助理教授陳丹琦團隊又有了新論文了。

近期，諸如「長思維鏈」等技術的興起，帶來了需要模型生成數萬個 token 的全新工作負載。

大多數語言模型都基于 Transformer 架構，其在進行自回歸解碼（即逐字生成文本）時，需要將所有先前 token 的注意力狀態存儲在一個名為 KV 緩存的內存區域中。

Transformers KV Caching Explained | by Jo?o Lages | Medium

KV 緩存是模型進行快速推理的基石，但它的大小會隨著輸入文本的長度線性增長。例如，使用 Llama-3-70B 模型處理一個長度為 128K token 的提示（這大約相當于 Llama 3 技術報告本身的長度），就需要分配高達 42GB 的內存專門用于存儲 KV 緩存。

許多先前的工作意識到了這個問題，并提出了從內存中丟棄（驅逐）部分鍵值對的方法，以實現所謂的「稀疏注意力」。然而，在一個公平的環境下對它們進行橫向比較卻異常困難。

生成過程 = 預填充（對輸入進行前向傳播并保存鍵值對）+ 后填充（一次解碼一個輸出詞元）。

有些論文旨在加速預填充階段；另一些則忽略該階段，轉而致力于最小化后填充階段的內存開銷。同樣，有的研究側重于吞吐量，而另一些則著力于優化內存使用。

陳丹琦團隊提出了「KV 足跡」作為一種統一的度量標準，它是在所有時間步中，未被逐出的鍵值緩存條目所占比例的聚合值。這一個指標就同時涵蓋了預填充和解碼兩個階段的全部開銷，使得在同等基礎上比較不同方法成為可能。

論文標題：Cache Me If You Can: How ManyKVsDoYouNeed for Effective Long-Context LMs?
論文地址：https://arxiv.org/pdf/2506.17121v1
代碼地址： https://github.com/princeton-pli/PruLong

為了確保比較的實用價值，團隊定義了「關鍵 KV 足跡」：即在模型性能相對于完整的全注意力機制不低于 90% 的前提下，一個方法所能達到的最小 KV 足跡。這個「90% 性能」的硬性標準，確保了我們比較的是真正有用的、未嚴重犧牲模型能力的優化方法。

該度量標準揭示了先前 KV 驅逐方法存在的高峰值內存問題。其中后填充驅逐由于與預填充階段的驅逐不兼容，導致其 KV 足跡非常高。團隊對這類方法進行了改進，使其能夠在預填充期間驅逐 KV，從而顯著降低了 KV 足跡。

接著，團隊轉向「新近度驅逐」方法，并在此基礎上提出了 PruLong，這是一種端到端的優化方法，用于學習哪些注意力頭需要保留完整的 KV 緩存，而哪些則不需要。PruLong 在節省內存的同時保持了長上下文性能，其 KV 足跡比先前的方法小 12%，并且在具有挑戰性的召回任務中保持了原有的性能。

KV 緩存驅逐的統一框架

測量關鍵的 KV 占用空間

給定一個包含個 token 的提示語，基于 Transformer 的語言模型通常分兩個階段來生成一個響應：

預填充

整個輸入序列在一次前向傳播過程中被處理。每個注意力頭 h 的鍵值狀態

被存儲在 KV 緩存中。這里，其中 d 是該注意力頭的維度。

解碼

逐個解碼生成這些 token，每次生成時都會讀取并更新 KV 緩存。

KV 緩存的存儲消耗會隨著提示長度和生成長度的增加而線性增長，研究人員提出了許多方法來解決這一開銷問題。總體而言，這些方法通過稀疏化注意力模式，從而允許某些 KV 條目被驅逐。

然而，這些方法針對推理流程的不同階段進行了定制：有些方法在預填充階段之后丟棄 KV 條目，而另一些方法則在預填充階段也對 KV 緩存進行修剪。這使得對不同方法進行公平且全面的比較變得困難。首先探討為何常用的 KV 緩存大小指標無法衡量模型在實際應用中的實用性。

在實際應用中，對長上下文進行單次前向傳播的預填充操作成本高昂。對于長輸入序列，將輸入序列分割成多個塊，并在多次前向傳播中處理這些塊的分塊預填充方法正日益成為標準實踐。這種方法通常能夠減少與長輸入相關的峰值 GPU 內存占用，并使得較短提示的解碼過程能夠與較長提示的額外塊同時進行。

此外，像多輪對話或交錯工具調用等場景，還需要多個解碼和預填充階段，這就需要一種全面的方法來衡量 KV 占用空間。而推測性解碼進一步模糊了預填充階段和解碼階段之間的界限，因為解碼過程變得更加依賴計算資源。

在考慮預填充和解碼過程中都進行多次前向傳播的推理情況時，「KV 占用空間」應考慮隨時間變化的內存使用情況。例如，它應反映出在分塊預填充過程中，是否在預填充完成之前就已經驅逐了 KV 條目。

具體的推理過程由輸入長度、輸出長度以及因方法而異的實現細節來表征。由于缺乏能夠捕捉所有這些細微差別的指標，本研究提出了一種理想化的指標，該指標能夠：（1）跟蹤整個預填充和解碼過程中的 KV 緩存內存使用情況；（2）考慮每個 KV 條目的生命周期，從而實現對不同方法的公平且全面的比較。

本研究檢查這些方法的注意力模式（圖 1），并將每個鍵值（KV）條目分類為：活躍的（在當前步驟中使用）、非活躍的（在當前步驟中存儲但未使用）或被驅逐的（在任何未來的步驟中都未使用，并從內存中移除）。本研究將 KV 占用空間定義為所有時間步中未被驅逐的注意力條目的數量。該數值被歸一化為完全因果注意力。

例如，在圖 1 中，KV 占用空間為 26/36=72.2%。一種理想的方法會盡早驅逐 KV，以盡量減少占用空間。本研究考慮了另一種指標，該指標跟蹤注意力矩陣中的峰值 KV 占用率。在實驗中，這兩種指標得出的結論相似。

本研究還討論了方法與實際性能指標（如總令牌吞吐量和 GPU 內存利用率）之間的關系。研究發現，在許多情況下，KV 占用空間與吞吐量密切相關，但具體的排名取決于 KV 驅逐之外的實現細節——不同方法在不同實現框架下的實際效率差異很大。

關鍵 KV 占用空間：以往的研究通常在固定的稀疏度水平下報告任務性能，但本研究認為，一個更有意義的指標是在保留大部分原始性能的情況下所能達到的稀疏度。本研究將關鍵 KV 占用空間定義為一種方法在長上下文任務中保留完整注意力性能的一部分（本文中）時所需的最小占用空間。低于此閾值，性能下降可能會過于嚴重，導致該方法無法繼續使用。

高效長上下文推理的現有方法

本研究調研了高效的長上下文方法，并討論了它們如何契合本研究的 KV 占用空間框架。表 1 概述了主要方法，展示了這些方法如何進行不同的權衡以及使用不同的稀疏性概念。

動態和預填充稀疏性方面：Native Sparse Attention、MoBA、QUEST 和 TokenButler 將 KV 緩存視為兩級層次結構，僅將相關的注意力塊從高帶寬內存（HBM）加載到片上 SRAM 進行處理。像 MInference 和 FTP 這類技術，在預填充階段使用動態稀疏注意力來近似全注意力。動態稀疏性方法會產生更多非活躍的 KV，能夠提升吞吐量，但它們并未減少 KV 內存，因此這些方法與本研究的關注點正交。

近期性驅逐：先前的研究確定了流式注意力頭，這些注意力頭僅關注局部滑動窗口和一組初始的「匯聚令牌」。驅逐遠距離的鍵值（KV）條目會大幅減少 KV 占用空間（圖 2），因為在上下文長度增加時，KV 緩存的大小保持固定，并且這種方法可在預填充和解碼過程中應用。然而，近期性驅逐可能會「遺忘」相關的遠距離上下文，這促使 DuoAttention 和 MoA 僅將一部分注意力頭轉換為流式頭。作為 KV 緩存壓縮的有前景的候選方法，后續將更詳細地討論這些方法。

后填充驅逐：我們使用「后填充驅逐」這一術語來指代在預填充階段結束后從鍵值（KV）緩存中刪除令牌的方法。這些方法依賴于通常基于注意力分數的啟發式規則來識別上下文中最重要鍵值對。這些方法可以在預填充后大量修剪鍵值對，并在解碼過程中減少 KV 內存。然而，在具有長提示和短生成的推理場景中，由于所有 KV 條目在預填充期間都保存在內存中，這也會在驅逐前導致相當大的峰值內存，后填充驅逐只能實現有限的 KV 占用空間減少。

正交技術：量化通過降低 KV 緩存的精度而非基數來節省內存，并且可以與本文考慮的任何方法結合使用。另一個方向是在預訓練新語言模型之前設計內存高效的架構。這可能涉及在查詢或層之間重用 KV 狀態，降低鍵值維度，或者交錯全局和局部注意力層。其他方法是用循環層、線性注意力或狀態空間層替換 softmax 注意力。這些方法與 KV 驅逐正交。

PruLong：一種用于注意力頭專業化的端到端方法

本研究探討過：驅逐「陳舊」鍵值對（KVs）雖能顯著降低內存占用，但可能導致重要歷史信息的丟失。這一發現推動了后續研究工作，旨在識別哪些注意力頭關注全局上下文、哪些聚焦局部上下文，從而僅對局部注意力頭中的 KVs 執行驅逐操作。

DuoAttention 將注意力頭分為兩類：檢索頭，從整個上下文中召回相關信息；流式頭，僅關注最近的 token 和輸入序列開頭的少量「匯聚」token。DuoAttention 通過將注意力機制表示為流式注意力和全注意力的疊加，并通過參數化來學習注意力頭的類型。

其中，和分別遍歷 Transformer 的層和個注意力頭。掩碼通過原始模型和插值模型的最終隱藏狀態之間的重建損失進行訓練，并且通過正則化鼓勵掩碼稀疏化。DuoAttention 使用長上下文訓練數據，該數據由合成的大海撈針任務組成。在收斂時，通過將的底部設置為 0，其余設置為 1，獲得的注意力頭稀疏度。

MoA 是另一種使用自然文本的方法，但當序列長度超過 8K 個 token 時，由于需要顯式存儲完整的注意力矩陣，難以擴展。

雖然 DuoAttention 在實證中表現出色，但團隊發現了幾種進一步降低其關鍵 KV 占用空間的方法。團隊結合這些見解，設計出 PruLong（長程精簡注意力機制），一種用于 KV 驅逐的端到端方法。PruLong 像 DuoAttention 一樣將注意力頭分為兩類，但在訓練目標、參數化和訓練數據方面進行了創新。接下來將依次介紹這些內容。

下一個 token 預測損失

PruLong（長程精簡注意力機制）直接最小化混合注意力模型的下一個 token 預測損失，而非最后一個隱藏狀態的重建誤差，這與這些模型在文本生成中的使用方式更為契合。

針對注意力類型優化離散掩碼

DuoAttention 學習一個連續的門控變量，該變量易于優化，但沒有反映出在推理過程中會被四舍五入為 0 或 1，因此引入了訓練-測試差距。PruLong（長程精簡注意力機制）將視為從由參數化的伯努利分布中抽取的二進制掩碼，并通過來自剪枝文獻的既定方法——將伯努利分布重新參數化為硬實體隨機變量，實現端到端優化。最終目標如下