KV キャッシュと LLM 推論の最適化

なぜ「KV キャッシュ」が推論の主役なのか

LLM と Transformer で見たとおり、LLM は 自己回帰でトークンを1つずつ生成します。素朴に実装すると、新しいトークンを出すたびに「これまでの全トークン」を入力し直してアテンションを計算することになります。系列長 n のとき、各ステップで n 個分のアテンションを計算するので、全体では O(n²) に膨らみます。

ここで効くのが KV キャッシュ です。アテンションは入力トークンごとに Query・Key・Value の3つのベクトルを作りますが、過去トークンの Key と Value は、後続のステップでも値が変わりません。つまり一度計算すれば使い回せます。これをメモリに貯めておくのが KV キャッシュであり、現代の LLM 推論で速度とメモリの両面を支配する中心的な仕組みです。

アテンションのどこを再利用できるのか

復習すると、自己回帰のアテンションは概念的に次の式です（数式は描画されないのでテキストで示します）。

Attention(Q, K, V) = softmax(QKᵀ / √d_k) · V

• Q（Query）: いま処理中のトークンが「何を探すか」
• K（Key）: 各トークンが「自分は何者か」を示す見出し
• V（Value）: 各トークンが実際に運ぶ中身

生成の各ステップで新しい Query は1トークン分だけ作られます。一方 K と V は「これまでの全トークン分」が必要ですが、過去分は前ステップと同一です。だから過去の K・V をキャッシュし、毎ステップ追加するのは新トークン1つ分の K/V だけにできます。

ステップ t:
  1. 新トークン x_t から q_t, k_t, v_t を計算
  2. K_cache に k_t を追記、V_cache に v_t を追記   ← ここが「キャッシュ」
  3. attn = softmax(q_t · K_cacheᵀ / √d_k) · V_cache  ← 過去K/Vは再計算しない
  4. 次トークンを予測 → t+1 へ

この結果、1トークンの生成計算は O(n²) ではなく O(n) 相当（過去 n 件との内積）に下がります。生成が「文の頭から徐々に遅くなる」ように感じるのは、n が伸びてこの内積コストとキャッシュ参照量が増えるためです。

推論は2フェーズに分かれる：Prefill と Decode

KV キャッシュを理解する鍵が、推論を Prefill と Decode の2段階で捉えることです。性能特性がまったく異なります。

Decode が メモリ帯域律速なのが重要です。1トークン生成するだけのために、モデルの全パラメータと全 KV キャッシュを HBM から読み出す必要があり、演算量に対して読み出し量が大きい。だから Decode の高速化は「計算を速くする」より「読み出すデータ量を減らす／読み出し効率を上げる」方向が効きます。量子化や MQA/GQA が効くのはこのためです。

メモリ消費を見積もる：長文・大バッチで爆発する

KV キャッシュの泣きどころはメモリです。サイズは次のおおよその式で見積もれます。

KVキャッシュ量(バイト)
  ≒ 2 × 層数 × KVヘッド数 × head_dim × 系列長 × バッチサイズ × 精度バイト数
        ↑ KとVで2倍

具体例として、層数48・KVヘッド数(=注意ヘッド数)64・head_dim 128・FP16（2バイト）のモデルで、系列長4096・バッチ1なら概算は次のとおりです。

2 × 48 × 64 × 128 × 4096 × 1 × 2 バイト
  ≒ 6.4 GB  （1リクエストあたり、4Kコンテキストで）

たった1リクエストで数 GB です。系列長とバッチサイズに比例するため、同時実行ユーザーを増やしたり長文を扱ったりすると、モデル重みとは別枠で KV キャッシュが GPU メモリを圧迫します。最大同時実行数や扱える最大コンテキスト長は、しばしば計算速度ではなく KV キャッシュが置けるかで頭打ちになります。

KV を圧縮する：MHA → MQA → GQA

メモリ式の KVヘッド数 に注目すると、ここを削れば KV が直接縮むと分かります。これを狙うのが MQA / GQA です。

• MHA（Multi-Head Attention）: 従来型。Query ヘッドと同数だけ K/V ヘッドを持つ。表現力は高いが KV が最も大きい。
• MQA（Multi-Query Attention）: 全 Query ヘッドで K/V を1組だけ共有。KV を劇的に削減できるが、品質がやや落ちることがある。
• GQA（Grouped-Query Attention）: 中間。Query ヘッドをグループに分け、グループごとに K/V を共有。品質をほぼ保ったまま KV を数分の一に圧縮でき、現在の主流。

KV を縮める他の手段に、KV キャッシュ自体の量子化（FP16 → INT8/INT4 で1要素のバイト数を削る）や、古い／重要度の低いトークンの KV を捨てる手法（ウィンドウ化、エビクション）もあります。いずれも「メモリと品質のトレードオフをどこで取るか」という同じ問題の別解です。

PagedAttention：断片化を解いてスループットを稼ぐ

KV キャッシュを単純に「1リクエスト＝連続した大きなメモリ領域」として確保すると、深刻な無駄が生じます。生成がどこまで伸びるか事前に分からないため最大長で先に確保しがちで、実際には使われない領域（内部断片化）や、空きはあるのに連続領域が取れず割り当て不能になる外部断片化が起きます。

vLLM が導入した PagedAttention は、OS の仮想メモリ／ページングの発想を持ち込みます。

• KV キャッシュを**固定サイズの「ブロック（ページ）」**に分割する。
• 各リクエストの KV は、物理的に非連続なブロックの集まりとして管理し、論理→物理の対応表（ブロックテーブル）で参照する。
• 必要になった分だけブロックを割り当てるので、事前の過剰確保が不要になり断片化がほぼ消える。
• 共通プレフィックス（同じシステムプロンプト等）の KV ブロックを複数リクエストで共有でき、メモリをさらに節約できる。

従来:  [req A ████████____________]  ← 末尾は予約だけで未使用（内部断片化）
       [req B ██████______________]

Paged: ブロック単位で必要分だけ確保し、空きブロックを使い回す
       A: [b3][b7][b1]   B: [b2][b5]    共有: システムプロンプト→[b0]を両者で参照

これにより実効的に置けるリクエスト数が増え、GPU を遊ばせず詰め込めます。さらに、生成が終わったリクエストの枠へ次々と新規リクエストを差し込む 連続バッチング（continuous batching） と組み合わせると、スループットが大きく向上します。これが現代の LLM サービングエンジン（vLLM 等）の核心です。

周辺の最適化：FlashAttention・量子化・投機的デコード

KV キャッシュ以外にも、推論を支える定番があります。組み合わせて使われます。

特に FlashAttention は混同しやすいので区別します。FlashAttention は softmax の計算過程を効率化し、n×n の巨大なアテンション行列を丸ごと HBM に書き出さずに済ませる手法です。一方 KV キャッシュは K/V の再計算を省く手法。前者は計算の段取りの最適化、後者は結果の使い回しで、レイヤーが違うため両立します。

まとめ：KV キャッシュを軸に推論を捉える

LLM 推論のコストとレイテンシは、突き詰めると 「KV キャッシュをいかに小さく、いかに効率よく扱うか」 に集約されます。

この見立てを持つと、「コンテキストを伸ばすと急に重い」「同時接続を増やすと OOM する」「最初の1文字が遅い」といった現象が、Prefill / Decode と KV メモリの言葉で説明でき、打ち手の優先順位もつけられます。プロンプト設計で入力長を抑える工夫は プロンプトエンジニアリング や トークン化 と、長文を外部知識で補う設計は ファインチューニングと RAG と合わせて読むと、推論コストの全体像が線でつながります。