PagedAttentionとは?LLMサービングのKVキャッシュ断片化を減らし高スループット化する仕組みと使い道

PagedAttentionは、LLM推論で肥大化するKVキャッシュをページ単位で管理し、メモリ断片化と重複を抑える技術です。vLLMの中核になったこの仕組みを、課題、実装アイデア、実験結果、開発への応用まで日本語で整理します。

参考文献

Efficient Memory Management for Large Language Model Serving with PagedAttention

Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, Ion Stoica

論文を見る

今回の論文

今回取り上げるのは、Woosuk Kwon、Zhuohan Li、Siyuan Zhuang、Ying Sheng、Lianmin Zheng、Cody Hao Yu、Joseph E. Gonzalez、Hao Zhang、Ion Stoica による 2023 年の論文「Efficient Memory Management for Large Language Model Serving with PagedAttention」です。公開元は arXiv で、その後 SOSP 2023 でも発表されています。研究分野は LLM サービング、推論最適化、GPU メモリ管理です。URL は https://arxiv.org/abs/2309.06180 です。

この論文を選んだ理由は、いま多くの LLM プロダクトで使われている vLLM の核となる考え方を扱っており、モデルそのものではなく「どう配信基盤を作るか」で性能が大きく変わることを学べるからです。RAG、社内チャット、コード生成、エージェント実行基盤など、長いプロンプトや同時リクエストを扱うサービスほど実務への示唆が大きい論文です。

どんな技術か

PagedAttention は、LLM 推論時に使う KV キャッシュを、連続した大きな領域としてではなく、小さな固定長ブロックの集まりとして管理する技術です。発想としては OS の仮想メモリやページングに近く、論文ではこれを Attention の実行と一体で設計しています。

LLM の推論では、すでに読んだトークンの key と value を KV キャッシュとして保持しておくことで、次のトークン生成を高速化します。ただしこのキャッシュは、リクエストごとに長さが変わり、生成中にどんどん伸び、リクエスト終了とともに解放されます。ここを雑に確保すると GPU メモリが断片化し、実際には空きがあるのに大きな連続領域が取れず、バッチ数を増やせないという問題が起きます。

PagedAttention は、この動的で扱いにくい KV キャッシュをページ単位で切り分け、論理的には連続して見せつつ、物理的には非連続でも扱えるようにします。これにより、メモリ浪費を減らしながら、並列サンプリングやビームサーチのような複数系列生成でもキャッシュ共有をしやすくなります。単なるメモリ節約ではなく、LLM サービング全体のスループットを上げるための基盤技術です。

課題

この技術が解決しようとしている課題は、LLM サービングが計算性能だけではなく、KV キャッシュ管理に強く縛られていることです。特に生成系 LLM は、推論が 1 トークンずつ進むため、GPU の計算器よりもメモリ帯域やメモリ配置の悪さがボトルネックになりやすいです。

何が難しいのかというと、KV キャッシュの長さを事前に正確に決めにくいことです。ユーザーごとに入力長も出力長も異なり、同じモデルでも短い問い合わせと長い対話が混在します。そのため、各リクエストの最大長ぶんをまとめて確保すると、実際には使わない領域が大量に生まれます。

既存方式では、KV キャッシュを連続したメモリとして確保することが多く、ここに 2 つの限界がありました。1 つは内部断片化です。最大長を見込んで大きく予約したのに、途中で短く終われば未使用領域が無駄になります。もう 1 つは外部断片化です。長さの違うキャッシュを何度も確保と解放すると、細切れの空き領域ばかりが増え、次の大きな要求を入れにくくなります。

なぜこの課題を解く必要があるかというと、LLM サービスのスループットは「同時に何リクエストを GPU に載せられるか」で大きく決まるからです。論文では、既存システムでは実際に有効活用されている KV キャッシュ領域が 20.4% から 38.2% に留まるケースがあると報告しています。つまり、モデルが重いから遅いのではなく、メモリの使い方が悪いために同時実行数を増やせない場面がかなりあります。

実際の AI システムでは、この問題は長いプロンプトを使う RAG、候補を複数出すコード生成、複数ツール呼び出しを抱えるエージェント、同時接続数が多い社内チャットなどで顕在化しやすいです。モデル精度が十分でも、サービング基盤が詰まるとレイテンシもコストも悪化します。PagedAttention は、この「推論アルゴリズムではなくメモリ管理が律速になる」問題に直接切り込んだ技術です。

用語解説

KVキャッシュ
Transformer の自己注意で使う key と value を、すでに読んだトークン分だけ保存しておく仕組みです。LLM の生成では毎回過去全体を再計算すると遅すぎるため、このキャッシュが不可欠です。PagedAttention はこの KV キャッシュの保存方法そのものを作り替えています。
断片化(フラグメンテーション)
メモリに空きはあるのに、連続領域として取りづらくなって使いにくくなる現象です。この論文では、LLM サービングが遅い理由のかなりの部分を断片化で説明しており、単なる GPU 容量不足ではない点が重要です。
論理ブロックと物理ブロック
論理ブロックは各リクエストから見た「連続したトークン列」の単位で、物理ブロックは実際に GPU 上へ配置される固定長のメモリ単位です。PagedAttention はこの対応付けを表で持つことで、連続して見える KV キャッシュを非連続メモリに置けるようにしています。
ブロックテーブル
論理ブロックがどの物理ブロックに対応するかを管理する表です。OS のページテーブルに相当します。このテーブルがあることで、推論カーネルは必要な KV を順にたどって読めます。
Copy-on-Write
共有していたブロックに書き込みが必要になったときだけ複製する仕組みです。並列サンプリングやビームサーチでは、同じプロンプト部分を複数系列で共有できるため、PagedAttention の価値を理解するうえで重要です。

技術の仕組み

この論文のポイントは、Attention を速くする新しい数式ではなく、KV キャッシュを OS 的な発想で管理し、その前提で Attention カーネルを動かすことです。vLLM は PagedAttention の上に組まれたサービングエンジンであり、ブロック単位のメモリ管理とスケジューリングを組み合わせて性能を出します。

基本アイデア

従来は、1 リクエストの KV キャッシュを 1 本の連続領域として持つことが多く、途中で伸びるたびに大きな空きを予約したり、最初から最大長ぶん確保したりしていました。PagedAttention では、これを固定長ブロックの列として管理します。

各リクエストは、論理的には「先頭から末尾までつながった KV キャッシュ」を持っていますが、実際の GPU メモリでは、その中身があちこちの物理ブロックに分散していても構いません。重要なのは、どの論理ブロックがどの物理ブロックに対応するかをブロックテーブルで引けることです。

この設計により、必要なぶんだけ後ろにブロックを追加すればよくなります。最大長を最初から予約しないため、短い応答が多いワークロードで特に効きます。

モデル構造ではなくサービング基盤の変更

PagedAttention は新しい LLM ではありません。Transformer の学習方法や重みは変えず、推論時の KV キャッシュ配置と Attention の読み出し方を変える技術です。そのため既存モデルを比較的そのまま活かせる一方で、サービングランタイムやカーネル実装への理解が必要です。

論文では、この非連続な物理配置を前提にしても Attention が正しく計算できるよう、PagedAttention カーネル側でブロックテーブルを参照しながら KV を順に読みにいく方式を採っています。つまり、メモリマネージャと推論カーネルを別々に最適化するのではなく、最初から一体設計しています。

KV キャッシュマネージャ

vLLM の KV キャッシュマネージャは、GPU メモリ上に物理ブロックを大量に用意し、各リクエストには論理ブロック列を割り当てます。トークン生成が進むと末尾ブロックへ詰めていき、いっぱいになったら新しい物理ブロックを確保してブロックテーブルを更新します。

この方式の重要な点は、物理ブロックがすべて同じサイズなので、外部断片化をほぼ消しやすいことです。未使用部分が出るのは最後のブロックだけになりやすく、無駄が局所化されます。また、リクエスト終了時には、そのブロック群だけを返せばよいので、開放処理も単純です。

並列サンプリングやビームサーチでの共有

この論文が実務的に面白いのは、単に断片化を減らすだけでなく、共有も前提にしているところです。たとえば 1 つの入力プロンプトから複数候補を生成する並列サンプリングでは、分岐するまでは同じプロンプト部分の KV キャッシュを共有できます。

従来方式では、各候補系列のために同じプレフィックスを丸ごと複製しがちでした。PagedAttention では、共有ブロックを複数系列から参照し、分岐後に書き込みが発生した時点だけ Copy-on-Write で分けます。これにより、複数候補生成時のメモリ膨張をかなり抑えられます。

スケジューリングとの組み合わせ

LLM サービングでは、リクエスト単位ではなく反復単位でバッチを入れ替える iteration-level scheduling が重要です。vLLM はこの細かいスケジューリングと PagedAttention を組み合わせることで、各反復で終わった系列を外し、新しい系列を詰め込みやすくしています。

ここで効いてくるのが、KV キャッシュを柔軟に伸縮できることです。連続領域前提だと、新しい系列を入れたいのに適切な空きがなく、せっかく計算資源が空いても GPU に乗せられないことがあります。PagedAttention は、細かいブロック単位で詰め直せるため、この待ち時間を減らしやすいです。

分散実行にもつながる考え方

論文では vLLM を分散サービング基盤として設計しており、中央スケジューラが GPU ワーカーへ指示を出す形を取ります。ここでも PagedAttention の発想は有効で、各ワーカー上の KV キャッシュを柔軟に扱いやすくなります。

この点は、単一 GPU の最適化に留まらず、今の推論基盤設計にもつながります。大規模サービスでは「どの GPU にどの系列を載せるか」と「既存キャッシュをどれだけ活かせるか」が重要なので、ページング的な管理は後続研究にもつながりやすい土台です。

実験と結果

論文の実験は、単純なカーネル速度比較ではなく、実際の LLM サービング性能をどれだけ押し上げるかを確認しています。モデル、ワークロード、デコード方式を変えながら、メモリ効率とスループットの両方を見ているのが特徴です。

何を検証したのか

主な検証項目は 3 つあります。1 つ目は、既存の LLM サービングシステムに比べてスループットがどれだけ改善するかです。2 つ目は、メモリ断片化や重複がどれだけ減るかです。3 つ目は、並列サンプリングやビームサーチ、共有プレフィックスのような実務でよくある生成パターンでも有効かです。

使ったデータセットと評価設定

評価には OPT 13B、66B、175B と LLaMA 13B が使われています。ワークロードは ShareGPT と Alpaca のトークン長分布をもとに合成されており、実際の LLM サービスに近い入力長と出力長のばらつきを持たせています。論文では ShareGPT のほうが Alpaca より入力が平均 8.4 倍、出力が平均 5.8 倍長いとされており、長文ワークロードで差が出やすい設計です。

比較対象は FasterTransformer や Orca です。指標としては、同じレイテンシ条件でのスループット、メモリ使用効率、同時バッチ数などが見られています。

スループットは 2 倍から 4 倍改善

論文全体の結論として、vLLM は既存の最先端システムと比べて、同程度のレイテンシで 2 倍から 4 倍のスループット改善を示しています。これはモデルサイズが大きいほど、系列が長いほど、また複雑なデコードほど効きやすいという結果でした。

ここで重要なのは、単に 1 リクエストを速くしたのではなく、「同じ GPU でより多くのリクエストを同時に回せるようになった」ことです。プロダクト観点では、まさにコスト効率の改善につながります。

既存方式はメモリをかなり無駄にしていた

論文では、既存方式で実際にトークン状態の保存に使われている有効メモリが 20.4% から 38.2% しかないケースを示しています。残りは予約済み未使用領域や断片化によるロスです。

PagedAttention はここをほぼ解消する方向に効きます。最後のブロックに少し未使用領域が残ることはありますが、リクエスト全体を大きく先取りする必要がないため、全体として「空いているのに使えない」状態を減らせます。これは推論最適化の論文としてかなり実務寄りの成果です。

長い入力や複雑なデコードほど差が広がる

結果が示しているのは、PagedAttention の価値が単純な 1 本生成よりも、長文や複数候補生成で大きくなることです。長いプロンプトでは KV キャッシュ自体が大きくなるため、断片化の悪影響も大きくなります。また、並列サンプリングやビームサーチでは共有プレフィックスを再利用できるため、重複削減の効果がより強く出ます。

この結果から言えるのは、RAG やコード生成、エージェントのような「入力も長いし候補も複数欲しい」用途ほど、PagedAttention の恩恵が大きいということです。

何に使える?

PagedAttention は、長いコンテキストか高い同時実行数のどちらかを必要とする LLM アプリで広く使えます。特に「モデルを替える前に、いまの GPU でどこまで伸ばせるか」を考える場面で有効です。

RAG や社内検索の応答基盤

RAG では、検索結果のチャンクを複数まとめて LLM に渡すため、入力が長くなりやすいです。PagedAttention を使うと、長い入力が来ても KV キャッシュの扱いが安定しやすくなり、同時接続数を落としにくくなります。

特に社内検索や規程検索のように、複数候補文書をまとめて読む設計では有効です。retriever の絞り込み精度だけに頼らず、generator 側へやや厚めの文脈を渡す設計を取りやすくなります。

コード生成や複数候補出力

コード補完や修正提案では、1 回の入力から複数候補を出したいことがよくあります。このとき並列サンプリングを行うと、同じプレフィックスを何度も持つ必要があり、従来は KV キャッシュが膨れやすかったです。

PagedAttention は共有プレフィックスをブロック単位で再利用しやすいので、複数候補生成を比較的安く実装できます。開発支援ツールやコードレビュー支援との相性はかなり良いです。

AIエージェントの長い実行履歴管理

エージェント型アプリでは、ツール呼び出し履歴、観測結果、過去の思考ログなどを長く保持したくなります。こうした履歴はセッションごとの差が大きく、長さも読みにくいため、KV キャッシュ管理が崩れやすいです。

PagedAttention 自体はエージェントのアルゴリズムではありませんが、長い履歴を抱えた多セッション運用の土台として役立ちます。特に複数ユーザーが同時に走る業務自動化系のエージェントで効きやすいです。

推論 API や SaaS のコスト改善

外部向け API や SaaS では、GPU をどれだけ高稼働させられるかが粗利に直結します。PagedAttention によるスループット改善は、そのまま 1 GPU あたりの処理件数増加につながります。

ここでの価値は、モデル精度を変えずに改善できることです。新モデルの再学習や大規模評価をせずとも、推論基盤の改善だけでコスト体質を変えられる可能性があります。

開発や事業へのヒント

この論文から得られる一番大きなヒントは、LLM プロダクトの競争力はモデル選定だけでは決まらないということです。同じモデルでも、サービング基盤の設計次第で同時実行数、レイテンシ、原価が大きく変わります。

まず疑うべきはモデルではなくメモリ配置

LLM が遅いと、つい量子化や軽量モデル移行を先に考えがちです。しかし実際には、KV キャッシュの断片化や重複が主因で GPU を使い切れていないことがあります。もし長文入力や複数候補生成が多いなら、モデル変更より先にサービング層を見直す価値があります。

小規模プロダクトでも応用しやすい

PagedAttention の良いところは、巨大研究所だけの話ではない点です。自前でカーネルを書く必要はありませんが、vLLM のような実装を選ぶだけでも効果が出る可能性があります。小規模な社内ツールや PoC でも、同じ GPU で捌けるリクエスト数が増えれば十分意味があります。

プロンプト共有を前提にした機能設計がしやすくなる

この論文は、同じプレフィックスを多系列で共有する設計と相性が良いです。つまり、テンプレート付き生成、長いシステムプロンプトを使う業務アシスタント、同じ文書に対する複数質問などは、基盤側で効率化しやすいと考えられます。

プロダクト設計の観点では、「同じ文脈を何度も使う機能」をむしろ積極的に作りやすくなる、という見方もできます。これは検索、コーディング支援、ドキュメント QA などで活かしやすい発想です。

今後注目すべき方向性

今後の推論基盤では、単なるモデル圧縮だけでなく、プレフィックス共有、prefill と decode の分離、スケジューリング最適化、キャッシュ再利用が重要になっていくはずです。PagedAttention はその出発点の 1 つとして理解しておく価値があります。

特に長文 RAG やエージェントでは、モデルの能力だけでなく「長い文脈をどれだけ安く維持できるか」が重要になります。この論文は、そのための基本設計をかなり明快に示しています。

限界

PagedAttention にも限界はあります。まず、これは主に KV キャッシュ管理の改善であり、Transformer の計算量そのものを線形化する技術ではありません。系列が極端に長くなれば、Attention 自体の計算コストや帯域制約は依然として効いてきます。

また、導入効果はワークロード依存です。短い問い合わせばかりで、1 リクエストずつ低レイテンシ重視で処理する構成では、論文ほど大きな差が出ない可能性があります。逆に長文かつ高並列なサービスほど効果が出やすいです。

実装面でも、PagedAttention は単なる設定変更ではなく、ランタイムとカーネルの協調設計が必要です。vLLM のような実装を使えば恩恵を受けやすい一方、自前推論基盤へ組み込む場合は保守負荷や GPU ごとの最適化も考える必要があります。

さらに、この論文は 2023 年時点のサービング課題を対象にしており、その後は prefill 最適化、分離実行、分散キャッシュ再利用など、より複雑な課題も研究が進んでいます。したがって、PagedAttention だけで現代的な大規模推論基盤の最適解が完成するわけではありません。

よくある質問

Q. PagedAttention は Attention の新しい数式ですか?

A. いいえ、主眼は数式の刷新ではなく、KV キャッシュの持ち方と Attention 実行のしかたです。同じ Transformer を使いながら、メモリ管理とカーネル側の読み出し方法を変えて性能を上げます。

Q. FlashAttention とどう違うのですか?

A. FlashAttention は Attention 計算そのもののメモリアクセス最適化に強い技術で、PagedAttention は生成時の KV キャッシュ管理と共有に強い技術です。前者はカーネル内部の計算効率、後者はサービング時のメモリ配置と運用効率に主眼があります。

Q. どんなサービスで特に効きますか?

A. 長い入力を扱う RAG、複数候補を出すコード生成、同時接続が多いチャット、長い履歴を持つエージェントで効きやすいです。共通点は、KV キャッシュが大きくなりやすく、同時実行数も重要なことです。

Q. 小規模なチームでも恩恵はありますか?

A. あります。自前実装は重いですが、vLLM などの既存基盤を採用するだけでも、同じ GPU でより多くのリクエストを捌ける可能性があります。PoC 段階でもコストと体感速度の両方に効くことがあります。

Q. これだけで推論基盤は十分ですか?

A. それだけでは十分とは言えません。量子化、スケジューリング、prefill 最適化、複数 GPU 配置なども重要です。ただし、PagedAttention はその土台としてかなり重要で、特に KV キャッシュ由来の無駄を見直す第一歩になります。

今日の学び

この論文は、LLM サービングが遅く高コストになりやすい原因として、KV キャッシュの断片化と重複に注目しました。これに対して、OS のページングに似た発想で KV キャッシュを固定長ブロックとして管理し、非連続メモリ上でも正しく Attention を実行できるようにしました。

ここから得られるヒントは、AI プロダクトの性能改善はモデル刷新だけではないということです。長文 RAG、コード生成、エージェントのような実アプリでは、キャッシュ共有やメモリ配置の設計が、そのままスループット、原価、UX に効いてきます。

関連記事