vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention

要約

高スループットの LLM 推論には、GPU メモリの効率的な使用が不可欠です。
以前のシステムでは事前に KV キャッシュ用にメモリを予約していたため、内部の断片化により容量が無駄になってしまいました。
OS ベースの仮想メモリ システムからインスピレーションを得た vLLM は、KV キャッシュへの動的なメモリ割り当てを可能にする PagedAttend を提案しました。
このアプローチにより断片化が排除され、より大きなバッチ サイズでの高スループット LLM の提供が可能になります。
ただし、物理メモリを動的に割り当てることができるようにするために、PagesAttention は KV キャッシュのレイアウトを連続した仮想メモリから不連続な仮想メモリに変更します。
この変更では、ページングをサポートするようにカーネルを書き換える必要があり、メモリ マネージャーを実装するためにフレームワークを提供する必要があります。
したがって、PagedAttendance モデルは、ソフトウェアの複雑さ、移植性の問題、冗長性、非効率性をもたらします。
この論文では、動的な KV キャッシュ メモリ管理のための vAttend を提案します。
Pagedtention とは対照的に、vAttend は KV キャッシュを連続した仮想メモリに保持し、すでに存在するデマンド ページング用の下位システム サポートを活用して、オンデマンドの物理メモリ割り当てを可能にします。
したがって、vAttend は、アテンション カーネル開発者にページングを明示的にサポートする負担を軽減し、サービス提供フレームワークでのメモリ管理の再実装を回避します。
vAttend によって、さまざまなアテンション カーネルの変更されていない実装に対してシームレスな動的メモリ管理が可能になることを示します。
また、vtention は、vLLM よりも最大 1.97 倍高速にトークンを生成し、入力プロンプトの処理は、FlashAttendant および FlashInfer の PagedAttendant バリアントよりも最大 3.92 倍および 1.45 倍高速に行われます。

要約(オリジナル)

Efficient use of GPU memory is essential for high throughput LLM inference. Prior systems reserved memory for the KV-cache ahead-of-time, resulting in wasted capacity due to internal fragmentation. Inspired by OS-based virtual memory systems, vLLM proposed PagedAttention to enable dynamic memory allocation for KV-cache. This approach eliminates fragmentation, enabling high-throughput LLM serving with larger batch sizes. However, to be able to allocate physical memory dynamically, PagedAttention changes the layout of KV-cache from contiguous virtual memory to non-contiguous virtual memory. This change requires attention kernels to be rewritten to support paging, and serving framework to implement a memory manager. Thus, the PagedAttention model leads to software complexity, portability issues, redundancy and inefficiency. In this paper, we propose vAttention for dynamic KV-cache memory management. In contrast to PagedAttention, vAttention retains KV-cache in contiguous virtual memory and leverages low-level system support for demand paging, that already exists, to enable on-demand physical memory allocation. Thus, vAttention unburdens the attention kernel developer from having to explicitly support paging and avoids re-implementation of memory management in the serving framework. We show that vAttention enables seamless dynamic memory management for unchanged implementations of various attention kernels. vAttention also generates tokens up to 1.97x faster than vLLM, while processing input prompts up to 3.92x and 1.45x faster than the PagedAttention variants of FlashAttention and FlashInfer.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google