KVキャッシュとは?LLMが長い文脈で重くなる仕組みを整理

KVキャッシュとは?LLMが長い文脈で重くなる仕組みを整理

長いやり取りを続けていると、AI の返事がじわじわ重くなってくることがあります。

コンテキストウィンドウの上限にはまだ余裕があるのに、GPU のメモリだけが先に苦しくなります。

その裏側で効いているのが KV キャッシュという仕組みです。

KV キャッシュが何を抱えているのかが分かると、文脈が伸びるほど重くなる理由も自然に見えてきます。

KV キャッシュとは何か

LLM は文章を1トークンずつ予測していく自己回帰モデルです。

次の1トークンを出すたびに、それまでの全トークンに対して注意(attention)のスコアを計算する必要があります。

このスコア計算に使うのが、トークンごとに作られる Key と Value という2種類のベクトルです。

Hugging Face のドキュメントでも、自己回帰モデルは毎回同じ計算を繰り返してしまうため、その結果を保存して再利用するのが KV キャッシュだと説明されています(Cache strategies)。

つまり KV キャッシュは「過去のトークンぶんの計算結果の置き場」で、速度のためにメモリを差し出す仕組みなのです。

そもそも注意(attention)のスコアがどうやって文脈をつかむのかは、アテンションの記事で順を追って整理しました。

アテンションとは?Transformerが文脈を捉える仕組み

なぜ文脈が伸びるとメモリを食うのか

保存するのは、トークン1つにつき Key と Value のペアです。

文脈が伸びればトークンが増えるので、抱えるペアの数もそのまま増えていきます。

しかもこのペアは層ごとに独立して持つため、モデルの層数ぶんの掛け算になるのです。

7B クラスでよくある「32層・隠れ次元 4096」の構成を fp16 で見積もると、こうなります。

1トークンあたり = 2(KとV) x 32(層) x 4096(次元) x 2(byte)
              = 524,288 byte = 0.5 MB

  1,024 トークン :  0.50 GB
  4,096 トークン :  2.00 GB
  8,192 トークン :  4.00 GB
 32,768 トークン : 16.00 GB

モデルの重みは読み込んだ時点で固定なのに、KV キャッシュだけは会話が伸びるほど増え続けます。

Hugging Face のドキュメントも、KV キャッシュがメモリの大部分を占めて長文脈生成のボトルネックになりうると書いています。

長い会話で急にメモリ不足になるのは、この増え方が原因でした。

キャッシュを切ることもできる

Transformers では、生成時に use_cache=False を渡すとキャッシュを使わない動作になります。

# キャッシュを無効にして生成する(毎回すべて計算し直す)
model.generate(**inputs, max_new_tokens=20, use_cache=False)

ただし毎トークンで過去ぶんを全部計算し直すことになるので、実用的な速度は出ません。

裏を返せば、いま普通に使えている生成速度は KV キャッシュが支えているということですね。

生成そのものを速くする別の工夫としては、小さなモデルに先読みさせて本命のモデルがまとめて答え合わせする、投機的デコーディングもあります。

投機的デコーディングとは?LLM生成が速くなる仕組み

メモリを削る3つの方向

抱える量を減らす手段は、大きく3つに分かれます。

1つめはオフロードで、いま計算している層以外の KV を CPU 側へ逃がす方法です。

2つめは量子化で、KV そのものを低いビット数で持ちます。

3つめは固定サイズのキャッシュで、最大長を先に確保して torch.compile などの最適化を効かせる形です。

# オフロード(GPU メモリ優先)
model.generate(**inputs, cache_implementation="offloaded")

# 量子化キャッシュ(backend は quanto がデフォルト、hqq も選べる)
model.generate(**inputs, cache_implementation="quantized",
               cache_config={"nbits": 4, "backend": "quanto"})

# 固定サイズ(JIT 最適化と相性が良い)
model.generate(**inputs, cache_implementation="static")

どれも無料ではなく、オフロードはデータの往復ぶんスループットが落ちます。

量子化キャッシュも、文脈が短く GPU に余裕がある状況ではかえってレイテンシを悪くすることがあると案内されています。

固定サイズは長さのばらつく用途だと、確保した領域が無駄になりがちです。

どれを選んでも「速度とメモリのどちらを取るか」の取引になるわけです。

モデル側で増え方を止める作りもある

ここまでは「トークンが増えたぶんだけ増える」前提で書いてきました。

ただし、注意を直近の一定範囲に限る sliding window attention や、区切りごとに区切る chunked attention を使う層では話が変わります。

Hugging Face のドキュメントによると、この種の層ではキャッシュがウィンドウ幅やチャンク幅に達した時点で伸びるのを止めるそうです。

最大長を大きく指定しても、その層のキャッシュはウィンドウ幅より大きくならないと明記されています。

モデルによって長文脈時のメモリの伸び方が違うのは、こうした構造の差から来ています。

API 越しに使う人にも効いてくる話

ローカルでモデルを回さない人にとっても、KV キャッシュは無関係ではありません。

同じ長いプロンプトを何度も投げる場面では、共通部分の KV を先に作っておいて使い回すというやり方があります。

Hugging Face のドキュメントにも、共通のプロンプトでキャッシュを埋めてから複数の生成に使い回す例(prefix caching)が載っていました。

この発想が分かっていると、変わらない部分をプロンプトの先頭にまとめ、変わる部分を後ろに置くという組み立てが自然に出てきます。

毎回中身が変わる文字列を先頭に置くと、後ろの共通部分まで使い回せなくなります。

最後に

KV キャッシュは、速く答えるために過去の計算結果を抱え込む仕組みでした。

抱えるぶんメモリを使うので、文脈が伸びるほど重くなるのは構造上の必然です。

会話が重くなってきたら、履歴を削るのがいちばん素直な対処になります。

コンテキストウィンドウやトークンの話とあわせて読むと、輪郭がはっきりするでしょう。

コンテキストウィンドウとは?AIが指示を忘れる仕組みを整理
量子化とは?AIモデルを軽くする仕組みを整理

以上です。

コメントを残す

メールアドレスが公開されることはありません。 ※ が付いている欄は必須項目です

CAPTCHA