microsoft/vattention

何を解決するか

vAttentionは、LLMサービングシステムにおけるKVキャッシュのメモリ管理の課題に対処します。具体的には、動的メモリ割り当てをサポートするためにカスタムアテンションカーネルを再実装する必要がある問題（PagedAttentionが要求するもの）を解決し、プリフィル中心のワークロードにおけるパフォーマンスを向上させます。

動作方法

vAttentionはCUDA仮想メモリAPIを用いて、仮想メモリの割り当てと物理メモリの割り当てを分離します。これにより、物理メモリは必要に応じてオンデマンドで割り当てられつつ、KVキャッシュは仮想メモリ上で連続した状態を維持できます。メモリが連続的に見えるため、他のシステムで実装された複雑なユーザースペースのオンデマンドページングを必要とせずに、変更のないアテンションカーネルで使用できます。

対象ユーザー

NVIDIA GPU上の推論最適化およびメモリ管理に注力するLLMサービングシステムの開発者および研究者向けに設計されています。

主な特徴

• カーネル互換性: 変更のないアテンションカーネルに対して動的メモリ割り当てをサポート。
• 仮想メモリの分離: CUDA仮想メモリAPIを用いて仮想メモリと物理メモリの割り当てを分離。
• パフォーマンス: 特にプリフィル中心のワークロードにおいて、PagedAttentionよりもパフォーマンスが向上。
• 統合性: Sarathi-Serveと統合されており、FlashAttentionやFlashInferなどのバックエンドをサポート。
• カスタムドライバー対応: デフォルトの2MBを超えて64KB～256KBの小さなページサイズを可能にする、修正済みNVIDIA UVMドライバーを含む。