NVIDIA/kvpress

解決的問題

部署長上下文大型語言模型（LLMs）成本高昂，因為鍵值（KV）快取會隨著令牌數量線性增長。例如，使用 float16 處理 Llama 3.1-70B 的 100 萬個令牌，可能需要高達 330GB 的記憶體。KVPress 透過壓縮 KV 快取，降低這些記憶體需求並提升解碼速度。

工作原理

KVPress 實現了多種「壓縮方法」——在預填入階段或解碼階段定期對 KV 對進行剪枝或合併。它與 Hugging Face 的 transformers 庫整合，提供一個自訂的 KVPressTextGenerationPipeline 來處理分詞與聊天範本。

壓縮透過以下幾種策略實現：

• 基於評分的剪枝：使用重要性評分（例如，隨機、逆範數或注意力權重）來移除重要性最低的 KV 對。
• 結構化壓縮：維持特定的快取模式，例如在各層間維持金字塔形大小，或僅保留初始與最近的令牌（StreamingLLM）。
• 降維：基於通道注意力評分壓縮鍵的維度。
• 進階技術：使用 CUR 分解、上下文重建（KVzip），或將被移除的令牌合併到存活的鄰居中。

適用對象

致力於長上下文 LLM 效率研究的研發人員與開發者，需要一個標準化框架來實現、基準測試與部署 KV 快取壓縮方法。

主要亮點

• 豐富的函式庫：支援多種無需訓練的壓縮方法，包括 SnapKV、StreamingLLM 和 KVzip。
• 彈性流水線：提供自訂的 transformers 流水線，便於整合與評估。
• 雙階段壓縮：支援在預填入階段與解碼階段（透過 DecodingPress）進行壓縮。
• 量化支援：與 QuantizedCache 相容，進一步減少記憶體佔用。
• 基準測試工具：包含 CLI 與筆記本，可用於測量準確度、峰值記憶體使用量與總時間增益。