NVIDIA KVPress：長上下文 LLMs 的 KV 壓縮工具箱

NVIDIA 已推出 KVPress，這是一個 Python 工具箱，旨在壓縮大型語言模型 (LLMs) 中的鍵值 (KV) 快取，以實現長上下文視窗的高效處理。透過減少存儲中間注意力結果所帶來的記憶體開銷，KVPress 使模型能夠處理更長的序列——例如那些用於上下文檢索和延伸推理所需的序列——而無需大量的 GPU 記憶體。

KV 快取記憶體瓶頸

在自回歸 LLMs 中，文本生成是逐個 token 進行的。為了避免為每個新生成的 token 重新計算所有先前 token 的表示，模型使用 KV 快取 來存儲注意力層的鍵 (K) 和值 (V)。雖然這種方式優化了計算，但快取大小會隨著上下文長度線性增長。

KV 快取的記憶體消耗由以下公式決定：

$$ ext{Size} \left( ext{KV} ight) = 2 imes ext{precision} imes n_{layers} imes n_{heads} imes d imes n_{tokens}$$

對於類似 Llama 3-70B 在 bfloat16 精度下運行且擁有 1M token 上下文的模型，KV 快取單獨需要約 327.6 GB。加上模型權重所需的 140 GB，總記憶體需求約為 470 GB，這意味著 KV 快取佔總記憶體使用量的約 70%。

KVPress：模組化壓縮工具箱

KVPress 提供了一系列「presses」——這些壓縮算法會修剪較不重要的 KV 對以減少記憶體佔用。這些 presses 透過 forward hooks 整合到模型的注意力層中，並針對 預填充階段，在快取達到最大時進行壓縮。

可用的壓縮技術

• KnormPress: 根據最低的鍵值範數修剪 KV 對。
• SnapKVPress: 修剪與最近查詢的低注意力權重相關的 KV 對。
• ExpectedAttentionPress: KVPress 作者創造的新技術（未發布），根據未來查詢的最低預期注意力權重修剪 KV 對。

KVPress 設計為模組化，使研究人員能夠輕鬆用新方法擴展工具箱，並與其他節省記憶體的技術結合，例如 transformers 庫中提供的 KV 快取量化。

效能與記憶體影響

壓縮 KV 快取可降低峰值記憶體使用量並提升生成速度。對於在 bfloat16 精度下、上下文長度為 128k 的 Llama 3.1 8B 模型，使用 KVPress 並將壓縮比設為 50% 時，可獲得以下改進：

• 記憶體減少：峰值記憶體使用量從 45GB 下降至 37GB。
• 速度提升：在 A100 GPU 上，解碼速度從每秒 11 個 token 提升至每秒 17 個 token。

基準測試與精度權衡

KVPress 包含一個 CLI，用於在標準長上下文數據集（如 RULER、InfiniteBench 和 Loogle）上基準測試壓縮技術。

在 RULER 數據集（上下文長度為 4k）上的基準測試顯示，最有效的組合是 AdaKVPress 和 ExpectedAttentionPress。然而，該工具箱指出了一個關鍵的權衡：隨著壓縮比的增加，模型精度通常會下降。這凸顯了繼續研究壓縮算法的必要性，以最小化對模型性能的影響。