今回の論文
今回取り上げるのは、Ramya Prabhu、Ajay Nayak、Jayashree Mohan、Ramachandran Ramjee、Ashish Panwar による論文「vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention」です。2024 年 5 月に arXiv で公開され、論文中では ASPLOS 2025 掲載予定とされています。研究分野は LLM 推論基盤、GPU メモリ管理、システム最適化です。URL は https://doi.org/10.48550/arXiv.2405.04437 です。
この論文を選んだ理由は、モデルそのものではなく、LLM を本番で速く安く動かすための基盤レイヤーを正面から扱っているからです。RAG、エージェント、長文チャットのように文脈が長くなるアプリでは、KV キャッシュ管理の設計がそのままコストとスループットに効きます。vAttention はそのボトルネックに対して、かなり筋のよいシステム設計を出しています。
どんな技術か
vAttention は、LLM 推論で使う KV キャッシュを、仮想メモリ上では連続した大きな配列として見せたまま、物理メモリだけを必要になったタイミングで少しずつ割り当てる技術です。
従来の LLM サービングでは、KV キャッシュの断片化を防ぐために PagedAttention が広く使われています。PagedAttention は物理メモリを動的に割り当てられる一方で、KV キャッシュを仮想メモリ上でも非連続に扱うため、注意計算カーネル側が「ページ分割されたキャッシュ」を理解しなければなりません。
vAttention はそこを逆転させます。GPU カーネルから見える KV キャッシュはあくまで連続メモリのままにし、物理メモリの割り当てだけを CUDA の仮想メモリ管理 API に任せます。ひとことで言えば、「KV キャッシュのページングをアプリ側で再実装するのではなく、システムの仮想メモリ機構を活かして解く」技術です。
課題
この技術が解決しようとしているのは、LLM サービングにおける KV キャッシュ管理の難しさです。
まず難しいのは、各リクエストの KV キャッシュが 1 トークンずつゆっくり伸びることです。最大コンテキスト長まで一括確保すると、実際にはそこまで使わないリクエストが多いため、GPU メモリがかなり無駄になります。これが内部断片化です。
では動的に確保すればよいかというと、既存の PagedAttention 方式には別の限界があります。KV キャッシュが仮想メモリ上でも非連続になるので、Attention カーネルはブロックテーブルをたどりながらデータを読まなければなりません。その結果、GPU カーネルを書き換える必要があり、新しい高速カーネルをそのまま流用しづらくなります。
さらに、サービングフレームワーク側にも複雑さが増えます。どのリクエストの KV キャッシュがどのブロックにあるかを、アプリ側のメモリマネージャが管理しなければならないからです。これは OS が本来持っているアドレス変換の役割を、ユーザー空間で重複実装しているとも言えます。
なぜこの課題を解く必要があるのかというと、実際の AI システムでは長文化と多並列化が進んでいるからです。長文 RAG、複数ステップのエージェント、継続的なチャットセッションでは、重みより KV キャッシュが先に支配的になることがあります。ここをうまく扱えないと、GPU は空いているのにバッチを積めず、推論基盤のコスト効率が落ちます。
用語解説
- KVキャッシュ
- Transformer の自己注意で、過去トークンの Key と Value を保存して再利用する仕組みです。生成トークンが増えるほど大きくなり、長文推論や多並列処理では GPU メモリの主要ボトルネックになります。
- PagedAttention
- KV キャッシュを固定長ブロックに分け、必要な分だけ物理メモリを割り当てる方式です。断片化対策として有効ですが、キャッシュが仮想的にも非連続になり、Attention カーネルやサービング層が複雑になります。vAttention はここを置き換えようとしています。
- 仮想メモリと物理メモリ
- プログラムから見える連続アドレス空間が仮想メモリで、実際に GPU 上に確保される領域が物理メモリです。vAttention はこの 2 つを分離して扱い、見た目の連続性と実際の節約を両立します。
- 内部断片化
- 必要量より大きな領域を確保してしまい、未使用部分が無駄になる状態です。LLM では最大コンテキスト長ぶんを先取り確保すると起きやすく、バッチサイズを押し下げる原因になります。
- TTFT（Time To First Token）
- リクエスト投入から最初の 1 トークンが返るまでの時間です。ユーザー体感に直結する指標で、prefill カーネルやメモリ管理のオーバーヘッドが大きいと悪化します。vAttention はここも改善対象にしています。
技術の仕組み
vAttention の中心は、KV キャッシュの「見え方」と「実際の割り当て方」を分離することです。これにより、従来カーネルをそのまま使いやすくしながら、断片化も抑えます。
基本アイデア
通常の予約型アロケータでは、各リクエストの KV キャッシュ用に最大長ぶんの連続領域を先に押さえます。この方式はカーネル実装が単純ですが、使われない領域が多くなります。
一方で PagedAttention は、必要時に物理メモリを足していく代わりに、KV キャッシュが仮想アドレス上でもページ単位の寄せ集めになります。すると、Attention カーネルは単なる配列アクセスでは済まず、ページ対応版に作り替える必要が出ます。
vAttention は、仮想メモリは最初に大きく連続予約しつつ、物理メモリの割り当てだけを後ろにずらします。これで GPU カーネルから見る KV キャッシュは連続のままです。
モデル構造ではなくサービング層の工夫
この論文の提案はモデル構造の変更ではありません。Transformer 自体も、Attention の数式も変えません。変えるのは KV キャッシュの置き方です。
そのため、FlashAttention や FlashInfer のような既存の Attention カーネルを、原則としてそのまま再利用できます。これは研究実装より実務で重要です。新しい高速カーネルが出ても、PagedAttention 対応版の移植を待たずに取り込みやすいからです。
予約するのは仮想空間だけ
vAttention では、各レイヤーの K キャッシュと V キャッシュに対して、最大バッチサイズと最大コンテキスト長を支えられるだけの仮想領域を先に予約します。ただし、この時点では物理メモリはまだ張り付きません。
論文では、64bit システムの仮想空間は十分に広いので、見た目の大きな連続領域を先に押さえても、制約になるのは主に物理メモリ側だと説明しています。つまり、節約すべきなのは「見た目のサイズ」ではなく「実体があるページ」です。
実行時に step で必要ページだけ張る
サービングフレームワークは、新しいリクエストに reqId を割り当て、各イテレーションの直前に step API を呼んで、現在の系列長に必要な物理ページが十分にマップされているか確認します。
decode では 1 イテレーションで各リクエストにつき 1 トークンしか増えないので、新しく必要になるページは高々 1 枚です。prefill では最初に複数トークンをまとめて処理するため、そのぶん必要ページ数も多くなりますが、vAttention はそこも吸収します。
重要な工夫1: 小さいページサイズ
CUDA の仮想メモリ API は標準では 2MB 単位の物理ページ割り当てが中心です。しかし LLM 推論では、1 トークンあたりの KV キャッシュ増分はもっと小さいため、2MB 単位だとまた内部断片化が起きます。
そこで論文では、CUDA Unified Virtual Memory ドライバを修正し、64KB から 256KB の細かいページを扱えるようにしています。これによって、「動的割り当てなのにページが粗すぎて無駄が戻る」問題を抑えています。
重要な工夫2: 割り当て遅延を計算で隠す
物理ページの割り当てには CUDA ドライバとの往復が必要で、同期的に行うと待ち時間が発生します。vAttention はここも LLM 向けに最適化しています。
論文では、前の計算と次の割り当てをオーバーラップさせる、いくつかの割り当てを先回りして行う、解放を即時ではなく遅延させて再利用しやすくする、といった工夫を入れています。これにより、動的割り当てのレイテンシをユーザーに見えにくくしています。
処理の流れ
実際の流れはかなり明快です。起動時に vAttention を初期化し、KV キャッシュ用の仮想バッファを予約します。リクエスト到着時に reqId を払い出し、各ステップで系列長を渡して必要ページをマップします。終了時には free_reqid を呼び、ページは即時または遅延解放されます。
この設計の良さは、サービングフレームワーク側の API が比較的薄いことです。PagedAttention のようにブロックテーブルを毎回組み立ててカーネルへ渡す必要がありません。
実験と結果
論文では、vAttention が本当に移植しやすいのか、性能でも勝てるのか、割り当て帯域は十分かを検証しています。
何を検証したのか
主な検証は 3 つです。1 つ目は、PagedAttention 版のカーネルではなく、既存の vanilla カーネルをそのまま使っても高性能を出せるかです。2 つ目は、prefill の TTFT と decode のスループットがどう変わるかです。3 つ目は、動的物理メモリ割り当てが LLM 推論の要求帯域を満たせるかです。
どんなモデルや条件で測ったのか
論文では Yi-6B、Llama-3-8B、Yi-34B を使い、1 台または 2 台の NVIDIA A100 で評価しています。prefill では FlashAttention と FlashInfer の PagedAttention 版と vanilla 版を比較し、decode では vLLM や FlashAttention 系の構成とも比較しています。
prefill は TTFT がかなり改善
prefill の結果はかなりわかりやすいです。たとえば Yi-6B の 192K コンテキストでは、FlashAttention の Paged 版は TTFT が 266 秒、vAttention で vanilla カーネルを使うと 67.81 秒でした。FlashInfer でも 81.87 秒から 65.03 秒に短縮されています。
Llama-3-8B でも同じ傾向で、192K コンテキスト時に FlashAttention の Paged 版 157.33 秒に対して vAttention は 42.31 秒でした。長文入力の prefill で、PagedAttention 用カーネルのオーバーヘッドがかなり大きいことを示しています。
decode はスループット最大 1.23 倍改善
論文の要約では、vAttention は PagedAttention ベースの FlashAttention / FlashInfer 構成と比べて、LLM serving throughput を最大 1.23 倍改善すると報告しています。decode では特に、非連続な KV キャッシュをたどる追加命令が不要になるぶん、既存の高速カーネルをそのまま活かしやすいのが効いています。
論文中では、PagedAttention 版カーネルは命令数増加やブロックサイズ依存の悪化も観測されています。つまり vAttention の改善は、単なる実装の好みではなく、メモリアクセス経路自体を素直に戻した効果だと読めます。
物理メモリ割り当て帯域は十分高い
動的割り当ては遅いのではないか、という疑問に対しても結果があります。論文では、最小の 64KB ページでも GPU あたり 7.59GB/s の物理メモリ割り当て帯域を出せたとしています。
一方、decode 時に必要なメモリ割り当てレートは、著者らの測定では最大でも 600MB/s 程度でした。つまり必要量に対して 1 桁以上余裕があり、適切にオーバーラップすれば割り当て帯域がボトルネックになりにくいことが示されています。
粗い 2MB ページは無駄が大きい
ページサイズの重要性も定量的に示されています。たとえば Llama-3-8B の TP-1 では、2MB ページだと 1 ブロックが 1024 トークン相当になり、リクエストあたり最大 128MB の無駄が出うると報告されています。64KB ページなら最大 4MB です。
この結果から、vAttention の価値は「動的割り当てできること」だけでは足りず、「LLM 推論の粒度に合わせて小さく割れること」まで含めて成立しているとわかります。
何に使える？
vAttention は研究向けの技巧というより、長文対応や高密度サービングを行うプロダクト基盤にそのまま効く技術です。
長文RAG基盤
長文 RAG では、検索結果を多めに詰め込むほど prefill が重くなり、KV キャッシュも増えます。vAttention を入れると、長文コンテキストでも TTFT を抑えながら、KV キャッシュ断片化も避けやすくなります。
特に、社内文書検索や契約書レビューのように 32K 以上を日常的に扱うサービスでは、応答待ち時間と GPU 利用効率の両方に効く可能性があります。
AIエージェントの長いセッション管理
エージェントはツール実行結果や中間状態を長く保持しがちです。すると decode 中も KV キャッシュが膨らみ続けます。vAttention のような動的管理は、エージェントを長く走らせるほど効いてきます。
また、Attention カーネル側を特別対応させなくてよいので、新しい最適化カーネルが出たときの追従もしやすく、エージェント基盤の性能改善サイクルを速めやすいです。
マルチテナントな推論API
複数顧客のリクエストを 1 つの GPU クラスタでさばく API では、メモリ断片化がそのまま収益性に響きます。vAttention は、最大長ぶんの先取り確保を避けつつ、PagedAttention より実装を軽くできるので、サービングレイヤーの保守性と性能の両方に利点があります。
特に「vLLM 依存を減らしつつ、FlashAttention や FlashInfer の素の進化をそのまま取り込みたい」という基盤方針と相性がよいです。
推論基盤の自前実装や高速化検証
この論文は、推論最適化をするときに必ずしもモデル圧縮から入る必要はないことも示しています。自前のサービングレイヤーを持つチームなら、KV キャッシュ管理の再設計だけでかなりの改善余地があるかもしれません。
開発や事業へのヒント
この論文から得られるヒントは、モデルの知能そのものより、基盤の「詰まり方」を観察したほうが勝ち筋になる場面が多いということです。
カーネル最適化の足を引っ張る設計を減らす
自分で AI アプリを作るなら、サービング基盤が新しい最適化カーネルをそのまま使える設計になっているかを見直す価値があります。PagedAttention のようにカーネル実装へ強く侵食する設計は、一度ハマると将来の追従コストが増えます。
vAttention の発想は、性能改善のために下層実装を複雑化しすぎない、というアーキテクチャの教訓として使えます。
OSやランタイムの機能を再利用する
論文の本質は、需要があるたびにアプリ側で専用の仕組みを作るのではなく、既存の仮想メモリ機構を転用したことです。これは LLM 基盤以外にも応用できる考え方です。
たとえばジョブキュー、キャッシュ、再試行制御、スケジューリングでも、アプリ層で独自管理しすぎると最適化の自由度を失いやすいです。既存基盤で吸収できる責務はそちらへ寄せたほうが、保守性と速度が両立しやすいです。
小規模チームほど「移植不要」は価値がある
小規模なプロダクトでは、GPU カーネルの追従実装に工数を割けないことが多いです。vAttention のように既存カーネルを再利用しやすい方向は、少人数の開発でも扱いやすいです。
推論基盤を自前運用する SaaS では、モデルの差別化だけでなく、こうした運用レイヤーの簡素化が粗利や開発速度に効くことがあります。
今後注目すべき方向性
今後は、KV キャッシュ量子化、ページング、prefill/decode 分離、長文向け Attention 最適化が別々ではなく、まとめて設計される流れが強まりそうです。vAttention はその中で「メモリ管理はもっと OS 的に解ける」という方向を示しています。
この先、CUDA 側の細粒度ページ対応が進んだり、主要サービングフレームワークに同種の設計が入ったりすれば、PagedAttention 一択ではなくなる可能性があります。これは論文結果から見た将来予想です。
限界
vAttention にも注意点があります。まず、論文の実装は CUDA の低レベル仮想メモリ API や UVM ドライバ修正に踏み込んでおり、導入のハードルは低くありません。単なる Python ライブラリ差し替えで済む話ではないです。
また、性能改善は主に KV キャッシュ管理やカーネル移植性に由来するため、短文・小バッチ中心のワークロードでは効果が小さい可能性があります。長文 prefill や高並列 decode のように、メモリ管理が効いてくる条件で真価が出ます。
さらに、論文は A100 上での評価が中心です。他 GPU、他ドライバ、他フレームワークで同じ効果がどこまで出るかは追加検証が必要です。特に本番導入では、自社ワークロードで TTFT、トークンスループット、OOM 発生率を見るべきです。
加えて、64KB から 256KB の細粒度ページを扱う工夫は有効ですが、そのぶん実装や運用の複雑さは残ります。システム機能を使うことでアプリ複雑性を減らす提案ですが、基盤エンジニアリング自体が不要になるわけではありません。
よくある質問
Q. vAttention は PagedAttention の完全な上位互換ですか？
A. 発想としては置き換え候補ですが、完全な上位互換とまでは言い切れません。vAttention は仮想メモリ API や細粒度ページ対応に依存するため、導入環境や実装工数の条件があります。ただし、KV キャッシュを連続メモリとして扱える利点はかなり大きいです。
Q. どんな場面で一番効きますか？
A. 長文コンテキスト、複数同時リクエスト、高い TTFT が問題になる場面で効きやすいです。特に RAG やエージェントのように入力も出力も長くなりやすいサービスでは、KV キャッシュ管理の影響が大きくなります。
Q. モデルの精度は変わりますか？
A. この論文はメモリ管理の提案であり、モデルの重みや推論アルゴリズム自体は変えません。そのため、精度改善を狙う技術ではなく、同じモデル品質をより効率よく提供するための技術だと理解するのが適切です。
Q. 量子化や KV キャッシュ圧縮と一緒に使えますか？
A. 考え方としては併用しやすいです。vAttention はキャッシュの配置と割り当てを扱い、量子化はキャッシュの表現サイズを下げます。論文内で詳細な併用評価まではしていませんが、役割が違うので組み合わせ余地はあります。
Q. 小規模チームでも参考になりますか？
A. かなり参考になります。ドライバ改修まで自前でやるのは重いですが、「アプリ側で複雑な専用機構を持ちすぎない」「既存の高速カーネルをそのまま使える構成を優先する」という設計思想は、小規模チームほど重要です。
今日の学び
この論文は、LLM 推論で KV キャッシュ断片化をどう防ぐかという課題を扱っています。提案された vAttention は、KV キャッシュを仮想的には連続のまま保ち、物理メモリだけを需要に応じて割り当てることで、この問題を解こうとしました。
そこから得られるヒントは、AI プロダクトの性能改善はモデル変更だけでなく、サービング基盤の設計でも大きく進むということです。とくに長文 RAG やエージェントでは、KV キャッシュ管理の設計がそのまま体感速度とコスト構造に跳ね返ります。