Re Reference AI

技術

KVキャッシュとは

KV Cache / Key-Value Cache

自己回帰生成において、過去のKey・ValueベクトルをGPUメモリに保持し再利用する推論最適化の仕組み

推論最適化Attention

ひとことで言うと

AIが文章を生成する際、過去の計算結果を覚えておいて使い回すことで速度を上げる仕組み。

概要

KVキャッシュ(KV Cache)とは、Transformerベースの自己回帰型LLMが1トークンずつ生成する際、Attention計算で使うKey・Valueベクトルを一度計算した後に保持しておき、次のトークン生成時に再利用する仕組み。 自己回帰生成では、新しいトークンを生成するたびにそれまでの全トークンに対するAttentionを再計算すると無駄な計算が発生するが、過去のKey・Valueをキャッシュしておけば新しいトークン分だけを計算すればよくなり、生成スループットを改善できる。 一方でキャッシュのサイズは系列長・バッチサイズ・層数・ヘッド数に比例して増大するため、長いコンテキストや大きなバッチでの推論時にはGPUメモリの主要なボトルネックの1つになる。

背景

自己回帰型のTransformerは、生成済みのトークン列全体を毎回入力してAttentionを計算し直すと、生成するトークンが増えるほど計算量が無駄に増えていく。 KVキャッシュは、過去のトークンに対するKey・Valueベクトルが後続の生成でも変化しないという性質を利用し、計算結果を再利用することでこの無駄な再計算を省くために導入された。

アーキテクチャ

TransformerのAttention層では、各トークンに対してQuery・Key・Valueベクトルを計算する。 自己回帰生成中、生成済みトークンのKey・Valueベクトルはそれ以降変化しないため、これをGPUメモリ上にキャッシュしておき、新しいトークンについてはQuery・Key・Valueを計算し、キャッシュされた過去のKey・Valueと合わせてAttentionを計算する。

ワークフロー

プロンプトの全トークンに対しKey・Valueを計算しキャッシュ(プリフィル) → 新しいトークンを1つ生成 → その新トークンのKey・Valueのみを計算しキャッシュへ追加 → キャッシュ全体を使ってAttentionを計算し次のトークンを生成 → このステップを繰り返す。

コード例

Hugging Face Transformersでuse_cacheを指定して生成する

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

inputs = tokenizer("Once upon a time", return_tensors="pt")
output = model.generate(
    **inputs,
    use_cache=True,  # KVキャッシュを有効化(デフォルトで有効)
    max_new_tokens=50,
)

print(tokenizer.decode(output[0], skip_special_tokens=True))

利点

  • 自己回帰生成における冗長なAttention再計算を省け、生成速度を改善できる
  • 長い出力を生成するほど、キャッシュによる計算量削減の効果が大きくなる
  • ほぼ全ての推論フレームワーク・APIにおいて、デフォルトで有効になっている

欠点

  • キャッシュのサイズが系列長・バッチサイズ・層数・ヘッド数に比例して増大し、GPUメモリを圧迫する
  • 長いコンテキストや大きなバッチサイズでの推論では、KVキャッシュがメモリ使用量のボトルネックになりやすい
  • キャッシュを効率化する実装(グループ化されたクエリAttention等)や量子化には追加の工夫を要する

比較

  • Attention — KVキャッシュは、Attention計算で使うKey・Valueベクトルを再利用する最適化手法
  • コンテキストウィンドウ — コンテキストウィンドウが長いほど、KVキャッシュのサイズも大きくなる
  • 量子化 — KVキャッシュのメモリ使用量を抑えるため、キャッシュ自体を量子化する手法も使われる
  • 推論 — KVキャッシュは、推論(特にデコードフェーズ)を高速化する代表的な最適化技術
  • 投機的デコーディング — 投機的デコーディングは、KVキャッシュ等と組み合わせてさらに推論を最適化できる
  • バッチ推論 — バッチ推論では、リクエストごとに異なるKVキャッシュを同時に管理せねばならず、実装が複雑になる
  • vLLM — vLLMは、PagedAttentionによりKVキャッシュの管理を効率化する代表的な実装

関連用語

AttentionTransformerコンテキストウィンドウ量子化推論投機的デコーディングバッチ推論vLLMプロンプトキャッシュGQAMQAFlashAttentionMLA

よくある質問

KVキャッシュがないとどうなる？

新しいトークンを生成するたびに、それまでの全トークンに対するKey・Valueの計算をやり直すことになり、生成が進むほど計算量が増え、応答速度が低下する。

コンテキストウィンドウを長くするとKVキャッシュはどうなる？

キャッシュに保持するKey・Valueの量も比例して増えるため、必要なGPUメモリが増加する。 これが長いコンテキストウィンドウを持つモデルの実運用コストが高くなる一因になっている。

参考文献

関連Zenn記事