KVキャッシュをOSの仮想メモリのようにページ単位で管理し、メモリ断片化を解消してGPUメモリ利用効率を最大化するLLM推論技術。vLLMが初めて実装した。
PagedAttention は、LLMの推論時に発生するKVキャッシュのメモリ管理問題を、OSの仮想メモリ管理の概念を応用して解決した画期的な技術です。UC Berkeley の Kwon et al. が 2023 年に発表し、推論エンジン vLLM の中核アルゴリズムとして実装されました。従来のKVキャッシュ管理では最大シーケンス長分のメモリを事前確保する必要がありましたが、PagedAttention は動的にページを割り当てることで60〜80%のメモリ無駄を解消します。
従来のLLM推論エンジンでは、各リクエストに対してKVキャッシュを連続メモリ領域に確保します。この方式には3つの深刻な問題があります。
最大シーケンス長(例: 2048トークン)分のメモリを事前に確保しますが、実際の生成長は予測できません。平均生成長が512トークンなら、75%のメモリが無駄になります。
リクエストの完了・開始が繰り返されると、メモリ上に使えない隙間(フラグメント)が生じ、連続領域を確保できなくなります。
Beam Search や Parallel Sampling で複数の候補系列が同じプレフィックスを共有する場合でも、従来方式では各候補にKVキャッシュを複製する必要がありました。
PagedAttention はKVキャッシュを固定サイズのブロック(ページ)に分割し、ブロックテーブルで論理ブロックと物理ブロックのマッピングを管理します。
| 概念 | OS仮想メモリ | PagedAttention |
|---|---|---|
| ページ | 4KB固定ブロック | KVブロック(16トークン等) |
| ページテーブル | 論理→物理アドレス変換 | ブロックテーブル |
| デマンドページング | アクセス時に物理ページ割当 | トークン生成時にブロック割当 |
| Copy-on-Write | 書き込み時のみコピー | 分岐時のみKVブロック複製 |
同一のシステムプロンプトを使う複数リクエストは、プレフィックス部分のKVキャッシュを物理ブロックレベルで共有できます。Copy-on-Write 方式で、分岐点以降のブロックのみを新規割り当てします。これにより、チャットボットのようなシステムプロンプト固定の用途でメモリ使用量とprefill計算を大幅に削減できます。
vLLM の論文では、PagedAttention により以下の改善が報告されています。
A1: vLLMのデフォルトは16トークンです。ブロックサイズを小さくするとメモリ効率は上がりますが、ブロックテーブルの管理オーバーヘッドが増えます。GPUのメモリアクセスパターンとの兼ね合いで、16が実用的なバランスとされています。
A2: はい。vLLM 0.3以降ではFlashAttention-2とPagedAttentionを組み合わせた FlashAttention with Paged KV Cache が実装されています。メモリ効率(PagedAttention)と計算効率(FlashAttention)の両方の恩恵を受けられます。
A3: ブロック単位の間接アクセスにより、連続メモリアクセスと比較してカーネル実装が複雑になります。また、ブロックテーブルの管理にCPU側のオーバーヘッドが発生します。ただし、メモリ効率の改善によるバッチサイズ拡大効果の方が大きく、総合的にはスループットが向上します。