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

KVキャッシュとは

KVキャッシュ（Key-Value Cache）は、Transformerベースの言語モデルが自己回帰的にトークンを生成する際、過去のトークンに対して計算済みのKeyテンソルとValueテンソルをGPUメモリ（VRAM）に保持し、次のトークン生成時に再計算を回避する最適化技術である。

Transformerの自己注意機構（Self-Attention）では、各トークン生成時に全ての過去トークンとの注意スコアを計算する必要がある。KVキャッシュなしの場合、N番目のトークンを生成するために1〜N-1番目の全トークンのKey/Valueを再計算する必要があり、全体の計算量はO(n²)となる。KVキャッシュによって過去のKey/Valueを再利用すれば、新しいトークンのKey/Valueのみ計算すれば良く、各ステップの計算量はO(1)、全体でO(n)に削減される。

KVキャッシュのメモリ消費

KVキャッシュのメモリ消費はモデルサイズ、シーケンス長、バッチサイズに比例する。

計算式: KVキャッシュサイズ = 2 × num_layers × num_heads × head_dim × seq_len × batch_size × bytes_per_element

この巨大なメモリ消費が、KVキャッシュの効率化が推論最適化の最重要テーマとなっている理由である。

KVキャッシュの効率化技術

PagedAttention（vLLM）

vLLM が導入した手法で、OSの仮想メモリ管理からインスピレーションを得ている。KVキャッシュを固定サイズのブロック（ページ）に分割し、物理メモリ上で非連続に配置する。従来の連続メモリ割り当てでは最大シーケンス長分のメモリを事前確保する必要があったが、PagedAttentionでは実際に使用した分だけページを割り当てるため、メモリ効率が大幅に向上する。メモリ無駄の削減率は実測で60-80%に達する。

GQA（Grouped Query Attention）

Key/Valueヘッドの数をQueryヘッドより少なくすることで、KVキャッシュサイズを削減する手法。Llama 2 70BがGQAを採用し、MHA（Multi-Head Attention）対比でKVキャッシュサイズを8分の1に圧縮した。推論品質への影響は最小限で、スループットは1.5-2倍に向上する。

MQA（Multi-Query Attention）

GQAの極端なケースで、全てのQueryヘッドが単一のKey/Valueヘッドを共有する。KVキャッシュサイズはヘッド数分の1に圧縮されるが、品質劣化がGQAより大きい。GPT-Jや初期のPaLMで採用された。

KVキャッシュ量子化

Key/Valueテンソルの精度をFP16からINT8やINT4に量子化することで、KVキャッシュサイズを2-4倍圧縮する。KIVI やQServeなどの手法がある。量子化による推論品質への影響は、INT8で無視できるレベル、INT4でもperplexityの増加は1-3%程度に抑えられることが報告されている。

Sliding Window Attention

Mistralが採用した手法で、過去の全トークンではなく直近W個のトークンのKVキャッシュのみを保持する。KVキャッシュサイズがシーケンス長に依存せずO(W)で一定となる。長距離依存関係は複数レイヤーの積み重ねでカバーする設計である。

KVキャッシュの運用上の課題

メモリフラグメンテーション

連続的なリクエスト処理で KVキャッシュの確保と解放を繰り返すと、GPUメモリが断片化してOOM（Out of Memory）が発生しやすくなる。PagedAttention やトークンバジェット管理で対策する。

プリフィル vs デコード分離

プロンプト処理（プリフィル）フェーズではKVキャッシュを一括生成し、トークン生成（デコード）フェーズではKVキャッシュを逐次追加する。この2フェーズの計算特性が大きく異なる（プリフィルは計算バウンド、デコードはメモリバウンド）ため、Splitwise やDistServeではプリフィルとデコードを異なるGPUに分離する手法が研究されている。

FAQ

KVキャッシュはユーザーが制御できますか？

API経由でLLMを利用する場合、KVキャッシュはサーバー側で自動管理されるため、直接制御する必要はない。ローカル推論（vLLM、llama.cpp等）を使う場合は、KVキャッシュのメモリ割り当て量やページサイズを設定パラメータで調整できる。vLLMでは gpu_memory_utilization パラメータでKVキャッシュに割り当てるVRAM比率を指定する。

KVキャッシュがVRAMを圧迫する場合の対処法は？

3つの対策がある。第一にGQAモデル（Llama 3等）を選択してKVキャッシュサイズ自体を削減する。第二にKVキャッシュ量子化（INT8/INT4）でメモリ使用量を2-4倍圧縮する。第三にvLLMのPagedAttentionを使ってメモリ効率を向上させる。これらを組み合わせれば、24GB VRAM（RTX 4090）でも70Bモデルの実用的な推論が可能になる。

Continuous BatchingとKVキャッシュの関係は？

Continuous Batching（連続バッチ処理）は、バッチ内の各リクエストが異なるシーケンス長を持つことを前提とし、完了したリクエストのKVキャッシュを即座に解放して新しいリクエストに再割り当てする。PagedAttentionとの組み合わせにより、GPUメモリの利用率を最大化しながらスループットを向上させる。vLLMはこの組み合わせにより、HuggingFace Transformersの素朴な実装比で2-24倍のスループット向上を達成した。