Zefan-Cai/KVCache-Factory

解決的問題

KVCache-Factory 解決了長上下文大型語言模型（LLM）推論過程中與鍵值（KV）快取相關的記憶體瓶頸。隨著上下文長度增加，KV 快取呈線性增長，消耗大量 GPU 記憶體並拖慢解碼速度。本專案提供一個統一框架，用於實現、比較和評估各種降低記憶體佔用的策略，同時不犧牲模型效能。

如何運作

此專案作為一個「實驗場」，將多種 KV 快取最佳化技術整合至單一評估介面之下。它支援以下幾類方法：

• 壓縮與淘汰：移除較不重要的 token（例如：StreamingLLM、H2O、SnapKV、PyramidKV）。
• 檢索：僅選擇與目前查詢最相關的 token（例如：Quest、L2Norm）。
• 合併：將多個 KV 對聚合為一個（例如：CAM）。
• 量化：降低快取值的精度（例如：KIVI、KVQuant、GEAR）。
• 卸載：將 KV 快取移至 CPU 記憶體，並逐個流式傳輸回 GPU，以在不導致 GPU 記憶體溢出的情況下維持完整精度（HeadInfer）。

包含用於 LongBench、RULER 和「針在 haystack 中」等標準基準的執行器，用於衡量這些方法對準確率和延遲的影響。

適用對象

本工具專為從事 LLM 效率研究的 AI 研究人員與開發者設計，尤其適合希望優化長上下文視窗效能或對比不同 KV 快取管理策略的人群。

亮點

• 統一介面：支援廣泛的基線方法，包括 PyramidKV、SnapKV、H2O 和 StreamingLLM。
• 多元化的最佳化策略：在一個地方整合了淘汰、檢索、合併、量化和無損卸載等多種策略。
• 全面的基準測試：內建支援 LongBench、RULER 和 Needle-in-a-haystack 評估。
• 硬體彈性：支援大模型（如 Llama-3-70B）的多 GPU 推論，並提供 FlashAttention v2 和 SDPA 兩種路徑。
• 可視化工具：包含用於可視化注意力模式的工具，幫助理解不同壓縮方法如何影響模型關注點。