KV Cacheとは？LLMが前のトークンを毎回計算し直さない仕組み
長い文章を生成するとき、LLMは次のトークンを1つ選ぶたびに、それまでの文脈を参照する。毎回、過去のトークンに対応するAttentionのKeyとValueまで計算し直すのは無駄が大きい。KV Cacheは計算済みのKeyとValueを保存し、次の生成ステップで再利用する仕組みだ。
「キャッシュがあれば文章をいくらでも長くできる」と思うかもしれない。しかし、保存にはメモリが要り、モデルの文脈長の制限も消えない。KV Cacheが主に節約するのは、過去の情報を繰り返し計算する手間である。
なぜKeyとValueを保存するのか
Self-Attentionでは、現在の位置のQueryと各位置のKeyを照合し、得られた重みでValueを集約する。因果マスクを使うデコーダーでは、過去の位置は未来のトークンを見ない。そのため、同じ入力の続きだけを生成する間、過去の位置で計算済みのKeyとValueは次のトークンが加わっても再利用できる。
たとえば入力が3トークンあり、4つ目を生成する場合を考える。最初に3トークンを処理して各層のKeyとValueを保存する。次のステップでは、新しいトークンのQuery・Key・Valueを計算し、そのQueryを「保存済みのKey 3個＋新しいKey 1個」と照合する。重みを使って「保存済みのValue 3個＋新しいValue 1個」を集約する。保存するのはQueryではなく、次のトークンから繰り返し参照されるKeyとValueだ。
実際にはこの操作をAttentionの層ごとに行う。Hugging Face Transformersの公式資料でも、キャッシュは層ごとにKeyとValueを持ち、新しい位置のKeyとValueを追加すると説明している。
最初の入力と、その後の生成は別の仕事
最初の入力文を処理してキャッシュを埋める段階をprefill、その後、1トークンずつ続きを作る段階をdecodeと呼ぶ。KV Cacheが特に効くのは後者だ。既存の文脈が長くなっても、過去の位置のKeyとValueを毎回作り直さずに済む。
ただし「1トークンの生成コストが一定になる」わけではない。新しいQueryは過去を含むKeyを参照するため、参照先が増えればAttentionの計算とデータ読み出しは増える。KV Cacheは過去の再計算をなくすが、長い文脈を読む仕事そのものはなくさない。
ここはTemperatureやTop-pとの違いでもある。TemperatureとTop-pは候補トークンの選び方を変える。KV Cacheは同じ生成処理を効率よく進めるための保存方法であり、選択規則そのものではない。
速さと引き換えに、メモリを使う
単純な全長保存型キャッシュの概算容量は次のように表せる。
は層数、 は同時処理する系列数、 はKey/Valueのヘッド数、 は保存するトークン数、 は1ヘッドの次元、 は1要素のバイト数だ。先頭の2はKeyとValueの2組を保存するために付く。
説明用に、32層・1系列・KVヘッド8個・4,096トークン・各ヘッド128次元・1要素2バイトとすると、 バイト、つまり512 MiBになる。これは仮定した形状のキャッシュ本体だけの概算だ。モデルの重み、途中の作業領域、実装上の余裕などは含めていない。
長い文脈や多くの同時リクエストでは、この容量が無視できなくなる。実際のモデルでは、KVヘッドを減らす方式、一定の窓だけを残すAttention、キャッシュの量子化やオフロードなどもあり、上の式をそのまま全モデルの実使用量とみなすことはできない。Hugging Faceのキャッシュ戦略資料も、速度とメモリ使用量のトレードオフを区別している。
キャッシュを再利用できる条件
保存したKeyとValueは、同じモデル設定と同じ入力の接頭部分に対応する計算結果だ。会話の途中を書き換えた、別の入力に差し替えた、あるいは位置の扱いを変えた場合に、以前のキャッシュをそのまま続きとして扱うと整合しない。複数の質問で共通の接頭部分を再利用する場合も、その部分が実際に同一であることが前提になる。
また、推論用の高速化を学習へそのまま持ち込まない。学習では過去位置を含む計算に勾配が必要で、Hugging Faceの説明もキャッシュを推論で使うものと位置づけている。
KV Cacheを理解する要点は、「過去のKeyとValueを保持することで再計算を減らす」「新しいトークンは保存済みの文脈を引き続き参照する」「保存するぶんメモリを使う」の3つだ。文章生成の速さと長い文脈のメモリ負担が同時に語られる理由は、ここにある。