Skip to content
Published on

vLLM 内部構造 (2) — PagedAttention はなぜ KV キャッシュをページに分けたのか

シェア
Authors

なぜ今ページの話を持ち出すのか

第1回でリクエストが通る経路を描いた。今回は、その経路の中でGPUメモリを最も多く消費する部分、すなわちKVキャッシュを見る。

生成中のリクエストは一つひとつ、これまでに見たすべてのトークンのキーと値のテンソルを保持していなければならない。これを捨ててしまうと、トークンをもう一つ作るたびに文全体を再計算することになるからだ。問題は、この塊がリクエストごとにサイズが異なり、生成が進むほど大きくなり続けるという点である。どれだけ大きくなるかは事前にはわからない。モデルがいつ生成を終えるかがわからないからだ。

サイズのわからないデータを継ぎ足し続けなければならない問題。コンピュータサイエンスはこの問題をすでに解いたことがある。オペレーティングシステムの仮想メモリとページングである。PagedAttentionの論文は、まさにその発想をアテンションに持ち込んだと自ら述べている。

内容は 2026-08-12 に公式ドキュメントとソースで確認しました。vLLM は変化が速いため、設定値と挙動は使用中のバージョンのドキュメントで再確認してください。

連続割り当てが生む二種類の無駄

ページング以前の方式を先に見ておくと、なぜ必要だったのかが見えてくる。素朴な実装では、リクエストが来ると「このリクエストが最大で使いうる長さ」の分だけ連続したメモリをあらかじめ確保する。安全で実装も簡単である。そして、非常に無駄が大きい。

最大2048トークンを想定し、リクエストごとに連続領域を事前予約した場合

リクエストA: [■■■■■□□□□□□□□□□□□□□□□□□□]  実使用320 / 予約2048
リクエストB: [■■□□□□□□□□□□□□□□□□□□□□□□]  実使用96  / 予約2048
リクエストC: [■■■■■■■■■■■■■■□□□□□□□□□□]  実使用1180 / 予約2048

■ 実際に使う領域   □ 確保したまま使わない領域(内部フラグメンテーション)

ここで生じる無駄は二種類ある。

一つ目は内部フラグメンテーションである。上の図の四角い空白すべてがこれに当たる。リクエストBは96トークンしか使わずに終わったのに、2048トークン分の領域を占有し続けていた。短い応答が多いサービスほどこの損失は大きくなる。しかも、この領域は他のリクエストが使うことはできない。すでにBの所有物だからである。

二つ目は外部フラグメンテーションである。リクエストが終わるたびにあちこちに空きの穴ができるが、それぞれの穴は小さく、新しいリクエストは大きな連続領域を要求する。空きメモリを全部合計すれば十分なはずなのに「メモリが足りない」という事態が起きるのは、こうして作られる。

ブロックテーブル: 論理的には連続、物理的には分散

PagedAttentionの解法は、リクエストに連続した広い領域を与えないことである。KVキャッシュをあらかじめ固定サイズのブロックに切っておき、リクエストには必要になるたびにブロックを一枚ずつ割り当てる。物理的にどこにあるかは関係ない。

では、アテンションカーネルは分散した断片をどうやって一つの文として見るのだろうか。その間を埋めるのがブロックテーブルである。リクエストごとに「自分の3番目のブロックは物理ブロック41番」という対応表を持っており、カーネルはこの表を見て必要なブロックを探しに行く。オペレーティングシステムがページテーブルで行っていることと同じである。

リクエストAが見る世界(論理) ブロックテーブル    実際のGPUメモリ(物理)
┌───────────────────────┐                       ┌────┬────┬────┬────┐
│ 論理ブロック 0 1 2 3  │   0 → 物理7           │ 0  │ 1  │ 2  │ 3  │
│ 一つの連続した文      │   1 → 物理2           ├────┼────┼────┼────┤
└───────────────────────┘   2 → 物理19          │ 4  │ 5  │ 6  │ 7  │
                            3 → 物理4           ├────┼────┼────┼────┤
                                                │ …  │ 19 │ …  │ …  │
                                                └────┴────┴────┴────┘

この構造は、先ほどの二種類の無駄を一度になくす。外部フラグメンテーションは定義上消える。すべてのブロックが同じサイズなので、どの空きブロックもどのリクエストにでも割り当てられるからである。内部フラグメンテーションは、最後のブロックに残る余りだけに縮小される。リクエスト一つあたり、最大でブロックサイズから1を引いた分だけである。論文がKVキャッシュの無駄をほぼなくしたと表現している根拠がこれである。

vLLMの設計ドキュメントは、このブロックがGPUスレッドブロックとは別の概念であるとわざわざ釘を刺している。ドキュメントを読むときに混同しやすい点なので、著者があらかじめ区別しておいたものである。

ブロックサイズのトレードオフ

では、ブロックサイズはどう決まるのか。vLLMには --block-size という引数がある。ただし、ソースの CacheConfig ではこの値はデフォルトが指定されない状態で宣言されており、ドキュメントにも、値を指定しなければデフォルト値を使うとしか書かれていない。実際に適用される数値はプラットフォームとアテンションバックエンドが決める。そのため、ここでは特定の数値をデフォルト値だと断定しない。デフォルト値はバージョンと環境によって異なるので、起動ログと使用中のバージョンのドキュメントを確認してほしい。

ただし、方向性ははっきりしている。

ブロックを大きく取ると、ブロック数が減って管理コストとテーブル参照が軽くなる。その代わり、最後のブロックで捨てられる領域が大きくなる。さらに、第5回で見るプレフィックスキャッシュは満杯のブロックしか再利用しないため、ブロックが大きいとプロンプトが少し違っただけで丸ごと外しやすくなる。

ブロックを小さく取ると無駄が減り、キャッシュヒットも細かい単位で合うようになる。その代わり、同じ文を表現するのに必要なブロック数が増え、ブロックテーブルも長くなる。

実務でこの値を最初にいじることはほとんどない。大抵は --gpu-memory-utilization と長さの上限を先に調整するほうが効果は大きい。ただし、プレフィックスキャッシュのヒット率が期待より低いとき、ブロックサイズが容疑者リストに挙がるということは覚えておく価値がある。

共有: 無駄を減らす二つ目の効果

ブロック単位で管理すると、もう一つ付いてくる利点がある。同じ内容を持つブロックを複数のリクエストが共に指し示せるということである。論文もこの柔軟な共有を、中心的な成果の一つとして並べて掲げている。

システムプロンプトが同じリクエストが百個同時に来る状況を思い浮かべればよい。連続割り当て方式では、同じ先頭部分を百回計算し、百セット分保存する。ブロック方式では、その先頭部分に当たるブロック群を一セットだけ置き、百個のブロックテーブルが同じ物理ブロックを指せばよい。このアイデアをリクエスト間に拡張したものがプレフィックスキャッシュであり、第5回のテーマである。

原論文が実際に主張していること

誇張を避けるため、出典を正確に記しておく。原論文はarXiv:2309.06180で、タイトルはEfficient Memory Management for Large Language Model Serving with PagedAttentionである。Woosuk Kwonらが書いたこの論文が、PagedAttentionとvLLMを同時に初めて提示した。

論文のアブストラクトが挙げている成果は、KVキャッシュのメモリの無駄をほぼなくしたこと、リクエストの内部・間で柔軟にKVキャッシュを共有すること、そして同水準のレイテンシで当時の最新システムと比べて処理量を2倍から4倍に高めたことである。

ここで注意すべき点がある。あの2〜4倍という数字は、2023年の論文が当時の比較対象と比べた数値である。今のあなたのデプロイで出る数字ではない。その間にvLLM自体も大きく書き直され、比較対象もすべて進化した。この論文は性能保証書としてではなく、設計の根拠として読むほうが正確である。

試してみる

参考資料