KVキャッシュ（ケーブイ キャッシュ）

KVキャッシュとは

KV キャッシュ（Key-Value Cache）は、Transformer ベースの LLM が自己回帰的にトークンを生成する際に、過去のトークンに対する Attention 計算の中間結果（Key テンソルと Value テンソル）をメモリに保持して再利用する最適化手法です。KV キャッシュがなければ、新しいトークンを 1 つ生成するたびに、過去のすべてのトークンの Key/Value を再計算する必要があり、計算量はシーケンス長の 2 乗に比例して増加します。

動作原理

Transformer の Self-Attention は、Query (Q)、Key (K)、Value (V) の 3 つのテンソルを使って計算されます。自己回帰生成では以下のフローで動作します。

Prefill フェーズ: 入力プロンプト全体を一度に処理し、各 Attention 層の K と V を計算してキャッシュに格納します。この段階は GPU の演算能力がボトルネック（compute-bound）です。

Decode フェーズ: 新しいトークンを 1 つずつ生成します。新トークンの Q を計算し、キャッシュ済みの K/V と Attention 計算を行います。新トークンの K/V もキャッシュに追加されます。この段階はメモリ帯域がボトルネック（memory-bound）です。

メモリ消費量

KV キャッシュは推論時のメモリ消費の主要因です。消費量は以下の式で概算できます。

KV キャッシュサイズ = 2 × レイヤー数 × ヘッド数 × ヘッド次元 × シーケンス長 × バッチサイズ × データ型サイズ

具体例として、Llama 3.1 70B（80 層、64 ヘッド、128 次元）でシーケンス長 128K、FP16 の場合:

2 × 80 × 64 × 128 × 131,072 × 1 × 2 bytes ≈ 330 GB

モデルの重み自体が約 140 GB（FP16）であるのに対し、KV キャッシュだけで 330 GB を消費する計算になり、長コンテキストでは KV キャッシュがメモリの支配的要因になります。

KV キャッシュの最適化手法

GQA（Grouped Query Attention）

複数の Query ヘッドで Key/Value ヘッドを共有する手法です。Llama 2 70B は GQA を採用し、8 つの KV ヘッドを 64 の Q ヘッドで共有することで、KV キャッシュサイズを 1/8 に削減しています。

MQA（Multi-Query Attention）

全 Query ヘッドで 1 つの Key/Value ヘッドを共有する極端な手法です。KV キャッシュサイズは最小になりますが、品質がやや低下する場合があります。

KV キャッシュ量子化

KV テンソルを FP16 から INT8 や INT4 に量子化してメモリ使用量を削減します。KV キャッシュの量子化は、モデルの重みの量子化よりも品質への影響が小さいとされており、INT8 量子化でほぼ無劣化、INT4 でも軽微な品質低下にとどまるケースが多いです。

PagedAttention（vLLM）

OS のページングと同様に、KV キャッシュを固定サイズのブロックに分割して非連続メモリに配置する手法です。vLLM で実装されており、メモリの断片化を防ぎ、同時バッチ処理数を大幅に増加させます。

FAQ

Q: KV キャッシュと Prompt Caching の違いは？

KV キャッシュは単一リクエスト内での Self-Attention 計算の再利用で、推論エンジンの内部最適化です。Prompt Caching は複数リクエスト間でのプレフィックスの再利用で、API レベルのサービス機能です。Prompt Caching の内部実装として KV キャッシュが使われています。

Q: KV キャッシュはなぜ Decode フェーズで特に重要ですか？

Decode フェーズでは 1 トークンずつ生成するため、新しい Q ベクトル 1 つに対して過去の全 K/V でAttention を計算します。KV キャッシュがなければ、過去トークンの K/V を毎回再計算する必要があり、生成速度が壊滅的に低下します。

Q: KV キャッシュのメモリが不足するとどうなりますか？

シーケンスの生成が途中で打ち切られるか、OOM（Out of Memory）エラーが発生します。推論エンジンによっては、古い KV エントリを追い出す eviction ポリシーや、KV キャッシュをメインメモリ（CPU RAM）にオフロードする機能で対処します。