microsoft/vattention

解決的問題

vAttention 解決了 LLM 服務系統中 KV 快取的記憶體管理挑戰。具體來說，它解決了需要重寫自訂注意力內核以支援動態記憶體配置（如 PagedAttention 所需）的問題，並提升預填入（prefill-bound）工作負載的效能。

工作原理

vAttention 使用 CUDA 虛擬記憶體 API 將虛擬記憶體配置與實體記憶體配置分離。這使得系統可以在需要時按需配置實體記憶體，同時保持 KV 快取在虛擬記憶體中的連續性。由於記憶體在虛擬位址空間中呈現連續，因此可直接用於未經修改的注意力內核，無需像其他系統那樣實作複雜的使用者空間需求分頁機制。

適用對象

專為 LLM 服務系統開發者以及專注於 NVIDIA GPU 上推論最佳化與記憶體管理的研究人員設計。

主要亮點

• 內核相容性：支援未經修改的注意力內核的動態記憶體配置。
• 虛擬記憶體分離：利用 CUDA 虛擬記憶體 API 分離虛擬與實體記憶體配置。
• 效能提升：在預填入密集型工作負載中，效能優於 PagedAttention。
• 整合支援：已整合至 Sarathi-Serve，並支援 FlashAttention 與 FlashInfer 等後端。
• 自訂驅動支援：包含修改後的 NVIDIA UVM 驅動，可啟用小於預設 2MB 的頁大小（64KB 至 256KB）。