OasisKV 把 LLM KV Cache 帶出 HBM 瓶頸

HBM 容量已成為一種稀缺且昂貴的資源,嚴重限制了推理批量大小和系統吞吐量。OasisKV 以記憶體為中心進行推理,透過解碼期間將完整的 KV 快取儲存與 HBM 解耦來減輕 HBM 容量壓力。

Hugging Face

OasisKV 來自 Microsoft Research、Imperial College London、KAIST 和 University of Edinburgh。是建基於 vLLM 的 LLM 推理系統設計。

長上下文及長篇推理會令 Large Language Model (LLM) 的 key-value (KV) cache 佔用大量記憶體和頻寬,High Bandwidth Memory (HBM) 容量遂成為批次大小與吞吐量的限制。OasisKV 不再把完整 KV cache 全部放在 HBM,而是在解碼階段只保留注意力較高的 KV 項目,其他資料存放於主機或遠端記憶體。

系統利用 speculative decoding (SD) 產生的 lookahead tokens,預測下一步可能重要的 KV blocks,再透過背景 attention pipeline 預先載入 HBM。論文報告指,在 2,048-token KV 預算下,準確度只比 full attention 低最多 0.7 分,推理工作負載吞吐量可達 dense vLLM 的 1.69 倍,多 GPU 長上下文服務最高達 2.1 倍。

  • 預取稀疏 KV,降低 HBM 容量壓力
  • 支援 host 或 remote memory 作較大容量層級
  • prefill–decode disaggregation 下吞吐量約為 dense 方法 2 倍
  • 每個請求所需 KV 可減少 6.5 至 9.7 倍
  • decode node host memory 可降低 2.2 至 2.6 倍

取捨在於系統依賴 lookahead 預測及稀疏 attention;預測失準可能影響準確度,且跨記憶體層級預取會增加系統複雜度。暫時只確認以 vLLM 實作,沒有提供 GGUF 檔案、量化版本、mmproj 附加檔案、Ollama 或 LM Studio 支援資料。

項目主頁 · Paper

Categories: 微軟, 推理引擎