Skip to content
Published on

KVキャッシュとPagedAttention — 推論メモリのすべて

🇰🇷KO🇺🇸EN🇯🇵JA
Split View필사 모드
Authors

はじめに

LLM推論において、GPUメモリは二つのものが分け合って使います。一つはモデルの重みであり、もう一つはKVキャッシュです。重みはモデルサイズが決まれば固定ですが、KVキャッシュは同時リクエスト数とシーケンス長に応じて際限なく膨らみます。そのため実務で「同時に何人を受け入れられるか」「どれだけ長いコンテキストを与えられるか」を決めるのは、ほとんどの場合KVキャッシュです。

この記事はKVキャッシュという一つのテーマを最後まで掘り下げます。まずKVキャッシュとは何で、なぜ必要なのか、それがなぜそれほど多くのメモリを消費するのかを計算で確認します。次に従来のメモリ管理の断片化の問題を見て、PagedAttentionがそれをどう解決するのか、prefix共有とKV量子化でどこまでさらに削れるのかを見ていきます。

核心原理: KVキャッシュとは何か

トランスフォーマーのアテンションは、各トークンに対してquery、key、valueという三つのベクトルを作ります。あるトークンを生成するとき、そのトークンのqueryは先行するすべてのトークンのkey、valueとアテンションを計算します。

attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V

ここで核心は、新しいトークンを生成するたびに、先行するトークンたちのkeyとvalueは変わらないという点です。100番目のトークンを生成するときに必要な1〜99番のトークンのK、Vは、すでに前のステップで計算したものとまったく同じです。であれば、毎回計算し直す理由はありません。

KVキャッシュはまさにこのK、Vを保存して再利用するものです。キャッシュがなければ、トークンを一つ作るたびに、これまでのすべてのトークンに対してK、Vを計算し直さなければなりません。シーケンスが長くなるほど、この重複計算は手に負えないほど大きくなります。KVキャッシュはこの重複を取り除き、decodeを実用的な速度にしてくれる必須の仕組みです。

KVキャッシュはなぜメモリを消費するのか

問題は、KVキャッシュがトークンごと、レイヤーごと、アテンションヘッドごとに積み重なるという点です。その大きさを形状で見ると次のようになります。

KVキャッシュ形状:
num_layers x 2 x batch x num_kv_heads x seq_len x head_dim

各項目の意味:
  num_layers   : トランスフォーマーのレイヤー数 (レイヤーごとにキャッシュ)
  2            : keyとvalueの二つ
  batch        : 同時処理するリクエスト数
  num_kv_heads : KVヘッド数 (GQA/MQAならqueryヘッドより少ない)
  seq_len      : 現在までのトークン数 (増え続ける)
  head_dim     : ヘッド一つの次元

総要素数に要素あたりのバイト数(例: FP16なら2バイト)を掛けると、メモリ使用量が出ます。

KVキャッシュバイト = num_layers * 2 * batch * num_kv_heads * seq_len * head_dim * dtype_bytes

メモリ計算の例

抽象的な式より数字のほうが実感できます。あるモデルが次のようだと仮定します(説明用の近似値)。

仮定:
  num_layers   = 32
  num_kv_heads = 8        (GQA適用)
  head_dim     = 128
  dtype_bytes  = 2        (FP16)

トークン1個あたり、リクエスト1個あたりのKVバイト:
  = 32 * 2 * 8 * 128 * 2
  = 131072 バイト (約128 KB)

シーケンス4096トークンのリクエスト1個:
  = 128 KB * 4096
  = 約512 MB

このようなリクエストを同時に32個受けると:
  = 512 MB * 32
  = 約16 GB (重みとは別に!)

つまりKVキャッシュだけで数十GBを軽く消費します。重みが16GBのモデルを80GB GPUに載せたからといって安心はできません。残りの空間をKVキャッシュがどれだけ効率よく使うかが、同時性とコンテキスト長を決めます。だからこそGQA/MQAでKVヘッド数を減らすことが、メモリの面で大きな意味を持ちます。上の式でnum_kv_headsが小さくなれば、そのままキャッシュが減るからです。

断片化の問題

KVキャッシュの大きさを減らすことと同じくらい重要なのが、その空間を無駄なく使うことです。従来の推論システムは、各リクエストに対して最大コンテキスト長ぶんの連続したメモリをあらかじめ確保していました。ところが実際の生成長は事前にわかりません。

最大2048トークンを確保したが、実際には100トークンしか生成しなかった場合:
[####... (100マス使用) ...______________ (1948マス無駄)]

リクエストごとにこのような無駄が積み重なると、GPUメモリは十分なのに
「これ以上受け入れる空きがない」という逆説的な状況が発生します。

これを内部断片化(internal fragmentation)と呼びます。さらに、リクエストが出入りすることでメモリに穴が空く外部断片化(external fragmentation)も発生します。大きな連続ブロックを確保しなければならないのに、散らばった空き空間しか残らず、割り当てに失敗するのです。従来の方式では、GPUメモリの相当な部分がこうしてただ捨てられていました。

PagedAttention: OS仮想メモリ式の管理

PagedAttentionの発想はオペレーティングシステムの仮想メモリから来ました。OSはプロセスに連続した仮想アドレス空間を見せますが、実際の物理メモリはページ単位で散らばっています。ページテーブルが両者をマッピングします。

PagedAttentionはKVキャッシュに同じことをします。KVキャッシュを小さな固定サイズのブロックに分け、トークンが積み重なるたびにブロックを一つずつ割り当てます。ブロックテーブルが各リクエストの論理的なトークン位置を物理的なブロックにマッピングします。

ブロックサイズ = 16トークンと仮定。

物理ブロックプール:
[blk0][blk1][blk2][blk3][blk4][blk5][blk6]...

リクエストA (40トークン):  ブロックテーブル -> [blk0, blk3, blk5]
                          (16+16+8トークン、最後のブロックだけ一部使用)
リクエストB (20トークン):  ブロックテーブル -> [blk1, blk2]

物理的に散らばっていても、論理的には連続のように動作します。

効果は明白です。あらかじめ最大値を確保せず、必要なぶんだけブロックを割り当てるので、内部断片化がほとんど消えます。ブロックが小さく均一なので、外部断片化も消えます。その結果、同じGPUメモリではるかに多くの同時リクエストを受け入れられます。無駄になっていた空間が、実際のスループットへ転換されるのです。

prefix共有とcopy-on-write

ブロック単位の管理のもう一つの大きな利点は共有です。複数のリクエストが同じ先頭部分(prefix)を持つとき、その部分のKVブロックを物理的に一組だけ置き、複数のリクエストがそれを指すようにできます。

共通のシステムプロンプトを持つ三つのリクエスト:

req1 ブロックテーブル: [共有blkS] -> [blk10] -> ...
req2 ブロックテーブル: [共有blkS] -> [blk11] -> ...
req3 ブロックテーブル: [共有blkS] -> [blk12] -> ...
                        ^^^^^^^
                  同じ物理ブロックを三つが共有

システムプロンプト、few-shotの例、マルチターン対話の共通履歴のように重なる部分が多いほど、この共有の利得が大きくなります。共有されたprefixはKV計算も一度だけで済みます。

ここで一つ注意すべき点があります。共有されたブロックを、あるリクエストが修正しなければならない場合です。たとえば分岐するシーケンスで片方だけトークンが変わると、そのブロックをコピーしてから修正してこそ、他のリクエストに影響を与えません。これがcopy-on-writeです。読むときは共有し、書く瞬間にコピーを作るのです。OS仮想メモリとまったく同じパターンです。

KV量子化: FP8とINT8

KVキャッシュがメモリのボトルネックなら、キャッシュ自体の精度を下げて大きさを減らす方法もあります。KVをFP16の代わりにFP8やINT8で保存すると、キャッシュメモリはおおよそ半分に減ります。

FP16 KV:  要素あたり2バイト  (基準)
FP8  KV:  要素あたり1バイト  (約50%削減)
INT8 KV:  要素あたり1バイト  (約50%削減、スケール/ゼロ点が必要)

decodeはメモリバウンドなので、KVの読み込み量が減ると速度にも役立ちます。ただし精度を下げると出力品質が微妙に低下することがあるため、モデルとタスクに応じて影響を検証する必要があります。一般的にKV量子化は重み量子化より品質への影響が小さい傾向ですが、無条件に安全だと断定はできません。重要なワークロードであれば、実際の評価指標で確認するのが安全です。

コンテキスト長とメモリのトレードオフ

KVキャッシュはseq_lenに線形に比例します。コンテキストを二倍に増やすと、リクエストあたりのKVキャッシュも二倍になります。つまり長いコンテキストと高い同時性は、同じメモリをめぐって競合します。

同じGPUメモリ予算の中で:
  コンテキスト4K  -> 同時リクエストを多く処理可能
  コンテキスト32K -> リクエストあたりKVが8倍 -> 同時リクエスト大幅減少
  コンテキスト128K -> リクエストあたりKVが爆増 -> 同時性が非常に制限される

サービングを設計するとき、「私たちは最大コンテキストをいくつまでサポートするのか」は、すなわち「同時に何人を受け入れられるか」と直結します。やみくもにコンテキストを大きくすると同時性が崩れます。このトレードオフを認識し、ワークロードに合わせて上限を定めることが重要です。

GQAとMQA: ヘッド数を減らしてKVを減らす

KVキャッシュ大きさの式でnum_kv_headsが掛けられているという点を、もう一度見てみましょう。もしKVヘッド数を減らせば、キャッシュは比例して減ります。これがGQA(Grouped-Query Attention)とMQA(Multi-Query Attention)の核心的な動機です。

MHA (Multi-Head Attention, 従来):
  queryヘッド32個、key/valueヘッドも32個
  -> KVキャッシュ = 32ヘッドぶん

MQA (Multi-Query Attention):
  queryヘッド32個、key/valueヘッドは1個 (すべてのqueryが共有)
  -> KVキャッシュ = 1ヘッドぶん (32倍削減!)

GQA (Grouped-Query Attention, 折衷):
  queryヘッド32個をグループにまとめ、グループごとにKVヘッド1個
  例) 8個のKVヘッド -> KVキャッシュ = 8ヘッドぶん (4倍削減)

MQAは最も積極的にKVを減らしますが、表現力がやや落ちて品質に影響が出ることがあります。GQAはその中間の折衷で、適度な数のKVヘッドを置いてメモリ削減と品質の両方を確保します。2026年現在、多くのオープンモデルがGQAをデフォルトとして採用していますが、これはサービングのメモリ効率がそれだけ重要になったからです。

KVキャッシュ削減効果 (head次元のみ比較):
  MHA(32 KVヘッド) : 基準100%
  GQA(8 KVヘッド)  : 約25%
  MQA(1 KVヘッド)  : 約3%

サービングエンジンを運用するとき、モデルがGQA/MQAを使っているかを知ることは、同時性の計画に直接的な影響を与えます。同じパラメータ数のモデルでも、KVヘッド構成によって受け入れられる同時リクエスト数が大きく変わるからです。

ブロックサイズという取っ手

PagedAttentionにおいてブロックサイズ(ブロック一つに収めるトークン数)は重要なチューニングの取っ手です。大きく取りすぎると最後のブロックの無駄(内部断片化)が大きくなり、小さく取りすぎるとブロックテーブル管理のオーバーヘッドとメタデータが増えます。

ブロックサイズが大きい場合 (例: 64トークン):
  - ブロックテーブルが短く管理が簡単
  - しかし最後のブロックの空き空間の無駄が大きい
    (あと3トークンだけ必要なのに64マスのブロックを丸ごと割り当て)

ブロックサイズが小さい場合 (例: 8トークン):
  - 最後のブロックの無駄が小さい
  - しかしブロック数が多くテーブル/メタデータが増加

実務: たいてい16前後の値がバランス点として使われる

ほとんどの場合はデフォルト値をそのまま使うのが無難ですが、非常に短いリクエストが多いワークロードと、非常に長いリクエストが多いワークロードでは、最適なブロックサイズが異なることがあります。これが「なぜブロック単位の管理が取っ手を提供するのか」を示す例です。

prefixキャッシュの寿命と無効化

prefix共有は強力ですが、共有されたブロックをいつまで生かしておくかが、また別の問題です。キャッシュは無限ではないので、ある時点では古いprefixブロックを空けて、新しいリクエストに空間を譲らなければなりません。

prefixキャッシュ管理の二つの軸:
  1) ヒット(hit): 新しいリクエストの先頭がキャッシュされたprefixと一致
     -> KV計算をスキップしてブロックを再利用 (利得)
  2) 追い出し(eviction): 空間が足りなければ古いprefixを削除
     -> たいていLRU(最も長く使われていないものから)方式

トレードオフ:
  prefixキャッシュを長く維持 -> ヒット率↑ だが利用可能メモリ↓
  早く追い出す -> 利用可能メモリ↑ だがヒット率↓

ワークロードが同じシステムプロンプトを使い続けるなら、そのprefixはほぼ常にヒットするので、維持する価値が大きいです。逆に、毎回まったく異なるprefixが入ってくるなら、キャッシュを維持するコストだけかかってヒットがありません。したがってprefixキャッシュは、ワークロードの性格がわかっているときに最もよく働きます。

メモリ予算を分ける方法

これまでの内容を総合すると、GPUメモリは結局いくつかの塊に分かれます。この配分を理解すると、サービングの設定が一目で見えてきます。

GPUメモリ予算の分解 (80 GB GPUの例):
  [モデル重み]          例: 16 GB (FP16 8B) 〜 40 GB (FP16 20B級)
  [活性化/作業バッファ]  例: 数GB (バッチ/シーケンスに比例)
  [KVキャッシュプール]   残り全部 -> 同時性とコンテキストを決定
  [安全余裕]            OOM防止用のヘッドルーム

核心: 重みを量子化で減らすほどKVキャッシュプールが大きくなり、
      それだけ多くの同時リクエスト、またはより長いコンテキストが可能。

この図が重要な理由は、サービング設定のほぼすべての決定が、この予算配分に帰着するからです。gpu-memory-utilizationを上げるということはKVキャッシュプールを大きくするという意味であり、量子化を有効にするということは重みを減らしてKVプールに空間を与えるという意味です。すべての取っ手は結局、この一枚の図の上で動きます。

落とし穴とトラブルシューティング

  • OOMが断続的に起きる: 平均的なリクエストではなく、同時に入ってきた長いリクエストたちがKVキャッシュを爆増させた場合です。最悪の同時性を仮定してメモリ予算を取りましょう。
  • メモリは余っているのに同時性が上がらない: 断片化、または保守的な事前割り当てが原因かもしれません。paged方式のエンジンを使っているか確認しましょう。
  • prefixキャッシュが期待ほど効かない: リクエストが実際には同じprefixを共有していないか、prefix長がブロックサイズより短く、共有単位が合わない可能性があります。
  • KV量子化後に品質が低下: タスクによって影響が異なります。量子化前後を同じ評価セットで比較しましょう。
  • GQAなのにメモリが減らない: 設定でKVヘッド数が実際に減っているか、キャッシュがKVヘッド基準で取られているか確認しましょう。

おわりに

KVキャッシュはLLM推論メモリの核心であり、同時性とコンテキスト長を左右する最も重要な変数です。その大きさはレイヤー、ヘッド、シーケンス長、精度の積で決まり、従来の管理方式は断片化で多くの空間を無駄にしていました。

PagedAttentionはOS仮想メモリの発想を借りてこの無駄をほぼ取り除き、prefix共有とcopy-on-writeで重複まで減らしました。KV量子化はさらに一歩進んで、キャッシュ自体を圧縮します。これらすべての技法の共通の目標は一つです。同じGPUでより多くのユーザーに、より長いコンテキストを提供すること。KVキャッシュを理解すると、サービングシステムの動作と限界がようやくはっきりと見えてきます。

参考資料

Discuss on Twitter • View on GitHub