PagedAttention(ページドアテンション)
KVキャッシュをOSの仮想メモリのようにページ単位で管理し、メモリ断片化を解消してGPUメモリ利用効率を最大化するLLM推論技術。vLLMが初めて実装した。
PagedAttentionとは
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の仕組み
PagedAttention はKVキャッシュを固定サイズのブロック(ページ)に分割し、ブロックテーブルで論理ブロックと物理ブロックのマッピングを管理します。
| 概念 | OS仮想メモリ | PagedAttention |
|---|---|---|
| ページ | 4KB固定ブロック | KVブロック(16トークン等) |
| ページテーブル | 論理→物理アドレス変換 | ブロックテーブル |
| デマンドページング | アクセス時に物理ページ割当 | トークン生成時にブロック割当 |
| Copy-on-Write | 書き込み時のみコピー | 分岐時のみKVブロック複製 |
ブロック割り当てフロー
- 新規リクエスト到着時、論理ブロック0を割り当て、空き物理ブロックにマッピング
- トークン生成が進み、現在のブロックが満杯になったら新規物理ブロックを割り当て
- リクエスト完了時、使用していた物理ブロックを解放プールに返却
- ブロックは不連続でよいため、外部断片化が原理的に発生しない
Prefix Cachingへの拡張
同一のシステムプロンプトを使う複数リクエストは、プレフィックス部分のKVキャッシュを物理ブロックレベルで共有できます。Copy-on-Write 方式で、分岐点以降のブロックのみを新規割り当てします。これにより、チャットボットのようなシステムプロンプト固定の用途でメモリ使用量とprefill計算を大幅に削減できます。
性能改善実績
vLLM の論文では、PagedAttention により以下の改善が報告されています。
- スループット: HuggingFace Transformers 比で最大24倍
- メモリ効率: 従来方式比でKVキャッシュメモリ使用量を最大55%削減
- バッチサイズ: 同一GPU上で同時処理可能なリクエスト数が2〜4倍に増加
FAQ
Q1: PagedAttentionのブロックサイズはどう決めますか?
A1: vLLMのデフォルトは16トークンです。ブロックサイズを小さくするとメモリ効率は上がりますが、ブロックテーブルの管理オーバーヘッドが増えます。GPUのメモリアクセスパターンとの兼ね合いで、16が実用的なバランスとされています。
Q2: PagedAttentionはFlashAttentionと併用できますか?
A2: はい。vLLM 0.3以降ではFlashAttention-2とPagedAttentionを組み合わせた FlashAttention with Paged KV Cache が実装されています。メモリ効率(PagedAttention)と計算効率(FlashAttention)の両方の恩恵を受けられます。
Q3: PagedAttentionの欠点はありますか?
A3: ブロック単位の間接アクセスにより、連続メモリアクセスと比較してカーネル実装が複雑になります。また、ブロックテーブルの管理にCPU側のオーバーヘッドが発生します。ただし、メモリ効率の改善によるバッチサイズ拡大効果の方が大きく、総合的にはスループットが向上します。