LowRAM:用預測壓縮長文本 KV Cache

長文本推理想慳顯示記憶體,不一定要保留每筆快取。LowRAM 用預測重建資料,但壓縮幅度與答案質素仍要取捨。

LowRAM: up to 26.9× KV-cache compression, retaining 94.5% of the uncompressed LongBench score on Mistral-7B prompts over

自回歸解碼依賴 KV 快取,其記憶體占用量會隨上下文長度呈線性成長,並限制批次大小與吞吐量。

現有的後訓練方法會利用詞元、層與特徵維度之間的冗餘,或降低數值精確度。LowRAM 找到一個互補的維度:在每個注意力頭內,值在很大程度上可由其對應的鍵預測。我們利用這一點採用值快取攤銷,在測試時僅需為每個注意力頭配備輕量級預測器,且只儲存其重建效果最差的詞元之殘差,即可達到近無損的 4 倍值壓縮。

預測器可透過元學習初始化,以大幅降低預填充開銷;也可在容許一次性成本的情境下從頭擬合,例如多查詢推理。我們將值攤銷與跨層低秩鍵壓縮結合,形成 LowRAM——一套非對稱壓縮管線,在三個模型系列中皆能維持接近無損的表現。LowRAM 在每個解碼步驟僅重建選定子集的 KV 配對,因此能將解碼開銷降至最低,並在 4 倍壓縮下實現最高 5 倍的峰值解碼吞吐量。

項目的 CompressedCache 可接入支援 Transformers 的模型流程,keys 與 values 能分別設定壓縮方式,不限於單純量化。

  • 跨層壓縮 keys,減少重複儲存。
  • 可配合量化及 token eviction,進一步縮小快取。
  • 提供預訓練預測器及 Llama、Mistral 範例。

測試前須準備 Python 3.12、CUDA GPU、Git LFS 及模型存取權限;預設 Llama-3.1-8B-Instruct。文件建議由 4K-token 範例開始,但提供的節錄未包含完整安裝及執行步驟;網站的 128K 展示則是錄製結果,並非即時推理。

Mistral-7B-Instruct-v0.3 在超過 10K tokens 的測試中,結合量化與淘汰可達 26.9 倍快取壓縮,保留未壓縮 LongBench 分數的 94.5%。這不是準確率,也不代表整體 GPU 記憶體縮小同樣倍數;團隊仍須以自己的文件檢查答案細節與速度。

項目主頁 · GitHub

Categories: 開源, 模型, 模型訓練, NVIDIA, 工具, LLaMa, Python