KVキャッシュは記憶のコスト — 文脈が伸びるほど decode が重くなる理由を実測する

Part 2 で、self-attention は系列長 T の二乗（O(T²)）で重くなることを実測しました。もし生成の 1 ステップごとに、この attention を毎回すべてのトークンに対して計算し直したら、文が長くなるたびに同じ二乗コストを何度も払うことになります。実際にはそんな無駄なことはしていません。過去に計算した Key と Value を キャッシュ しておいて、使い回しているからです。

では、その KV キャッシュは無料なのでしょうか。今回はそうではないことを実測で確かめます。キャッシュ自体が文脈の長さに比例して膨らみ続け、しかもその大きさを素朴に測ろうとすると、ある実装上の罠にはまって間違った数字を掴んでしまいます。この罠を実際に踏んでから抜け出し、最後には「decode が遅くなるのは計算ではなく読み出しのせいだ」という、この連載でいちばん実感しにくい主張を、自分の Mac で測った数字で裏付けます。

この連載の現在地 — キャッシュは無料ではない

• ここまで分かっていること — Part 1 で LLM は次トークン予測器で生成は逐次ループだと確かめ、Part 2 で self-attention が系列長の二乗（O(T²)）で重くなることを実測しました
• 今回新しく説明できること — その二乗の再計算を避けるための KV キャッシュの仕組み、キャッシュ自体が文脈長・層数・head 数に線形で膨らむこと、そして decode の速度が計算量ではなくキャッシュの読み出し量で決まる（memory-bound）こと
• 今回はやらないこと — サンプリングの工夫は Part 4、ゼロからの学習は Part 5、PHOTON のアーキテクチャの詳細は Part 6 へ送ります
• PHOTON のどの主張につながるか — 「キャッシュは文脈に線形で増える」ことと「そのキャッシュを減らす軸がある」ことの両方が、最終回「なぜ最大 475 倍が成立するのか」の材料の一つ（KV トラフィックの削減）になります。分解は最終回でまとめて行います

過去を毎回作り直すのをやめる — KV キャッシュの仕組み

自己回帰生成では、モデルは 1 トークン進むたびに、それまでのすべてのトークンに対して Key と Value を計算し直すことができます。ただし、これは無駄です。トークン 1〜5 の Key/Value は、6 番目のトークンを生成する時点でも 7 番目のトークンを生成する時点でも、まったく同じ値だからです。過去のトークンの埋め込みは変わらないので、そこから計算される Key/Value も変わりません¹。

そこで、一度計算した Key/Value を保存しておき、新しいトークン 1 個ぶんだけ計算して追加する、というのが KV キャッシュです。キャッシュなしでは各ステップの計算がそれまでの全トークン分の再計算（累積で二次的にスケール）になるところを、キャッシュありでは新規トークン 1 個分の計算だけで済みます¹。Part 1 の token_loop.py が「cached decode」で測っていたと注記していたのは、まさにこの仕組みのおかげです。

再計算そのものは避けられました。では、保存しておいた Key/Value はどこに消えるわけでもなく、メモリの上に居座り続けます。ここからが今回の本題です。

キャッシュのサイズは文脈に線形で効いてくる

KV キャッシュのサイズは、次の式で決まります。

KVキャッシュのバイト数 = 2 × 層数 × KVヘッド数 × head_dim × 文脈長 × dtypeのバイト数

先頭の 2 は、Key 用と Value 用の 2 本のテンソルを持つことに由来します²。層数・KVヘッド数・head_dim はモデルのアーキテクチャで決まる定数で、変わるのは 文脈長 だけです。つまりこの式は、文脈が伸びるほどキャッシュのサイズが 線形に 増え続けることを意味します。再計算は避けられても、置き場所は際限なく必要になるということです。

今回実験に使う mlx-community/Llama-3.2-1B-Instruct-4bit の実際の値を、モデルの中身から読み出してみます。

n_layers = 16, n_heads = 32, n_kv_heads = 8, head_dim = 64
KVキャッシュの dtype = float16（重みは 4bit 量子化でも KV は既定で fp16 のまま）

n_kv_heads が n_heads より小さいのには理由があります（後述の GQA）。まずは素直にこの式が実測と合うか確かめたいところですが、その前に 1 つ、実装上の落とし穴を踏んでおきます。

「キャッシュの大きさ」を素朴に読むと騙される

このモデルを動かしている mlx_lm ライブラリの実装を覗いてみます。KV キャッシュを作る make_prompt_cache() はモデルの層数ぶんの KVCache オブジェクトを返し、各オブジェクトは .nbytes というプロパティでキャッシュの大きさをバイト数で教えてくれます。素朴に考えれば、これを読めば理論式と一致するはずです。

ところが mlx_lm/models/cache.py のソースを読むと、KVCache は実データを直接持つのではなく、256 トークン単位でキャッシュ配列を事前確保 していることが分かります³。文脈が 1 トークン増えるたびに配列を作り直すのは非効率なので、あらかじめ 256 の倍数のサイズで確保しておき、実際に使っている長さは別途 offset という変数で覚えている、という設計です。

問題は .nbytes が 確保済みの配列そのもの（256 の倍数に切り上げたサイズ）を返すことです。実際に自分の Mac でこの境界をまたいで測ってみると、こうなりました。

文脈が 256 から 257 に、たった 1 トークン 増えただけです。それなのに .nbytes を素朴に読むと 8.0 MiB から 16.0 MiB へ、ちょうど 2 倍に跳ね上がります。一方、実際に使われているデータ量（offset でスライスした実データ、そして理論式)は 8.000 MiB から 8.031 MiB へ、ほんの少ししか増えていません。

これが今回踏んだ罠です。「キャッシュのメモリ使用量を測っている」つもりで、実は「ライブラリが先回りして確保しておいた容量」を測ってしまっていたわけです。事前確保は 256 トークン単位で効率よくメモリを再利用するための、ライブラリ側の合理的な最適化です。ただ、それを知らずに .nbytes の推移をそのままグラフにすると、実データの増加より遥かに大きい階段状のジャンプを「これがキャッシュの実サイズだ」と誤読してしまいます。理論式と正しく突き合わせるには、確保容量ではなく、実際に使われている長さ（offset）に基づく値を見る必要があります。

大きな文脈でも理論どおり線形 — 1K→32K の三点一致

罠を避ける方法が分かったので、今度は 256 の倍数ちょうどの文脈長を選んで測り直します。境界ぴったりであれば、確保容量・実データ・理論式の 3 つが完全に一致するはずです。

3 つの値は寸分違わず一致し、文脈が 2 倍になるとキャッシュもきっちり 2 倍になっています。2 × 16 × 8 × 64 × 文脈長 × 2 バイト という理論式が、実測とそのまま一致しました。文脈が 32,768 トークン（32K）まで伸びると、たった 16 層・1B パラメータの小型モデルでも KV キャッシュだけで 1 GiB に達します。これがモデルの重み（4bit 量子化でおよそ 0.7GB 程度）と別に必要になる、という点が、長い文脈を扱うときにメモリを圧迫する理由です。

decode が遅くなるのは「計算」ではなく「読み出し」

キャッシュのサイズが文脈に線形で増えることは分かりました。では、それは decode の速度にどう効いてくるのでしょうか。

直感的には、こう思うかもしれません。「新しく生成するのは 1 トークンだけなのだから、文脈が長くても短くても、1 ステップあたりの計算量は同じはずだ」と。実際、新しく計算する Q/K/V や出力の量は、文脈長によらず一定です。ところが、実測すると decode は明確に遅くなります。

文脈が 1,024 から 32,768（32 倍）に伸びると、decode 速度は 292 tok/s から 119 tok/s へ、約 2.45 倍遅くなりました。新しく計算する量は変わらないのに、です。

理由は、attention の計算そのものではなく、その計算に使う Key/Value を毎ステップメモリから読み出すコストにあります。新しいトークン 1 個の Query は、文脈中の すべて の Key/Value と照合されます。つまり、文脈が伸びるほど、1 ステップで読み出さなければならないデータ量も比例して増えるのです。計算する量ではなく、運んでくる量が効いてくるのです。この種のボトルネックを memory-bound（メモリ帯域律速）と呼びます²。文脈が短いうちはモデルの重み自体を読み出すコスト（文脈長によらず一定）が支配的なので速度低下は緩やかですが、文脈が伸びるほど KV キャッシュの読み出しコストの比重が増えていきます。

ここで 1 つ正直に注記しておきます。この「memory-bound」という整理は、GPU と別体のメモリ（VRAM）を持つ discrete GPU を念頭に語られることが多い話です²。今回測定に使った Apple Silicon（M5 Pro）は CPU と GPU がメモリを共有する unified memory アーキテクチャなので、事情がまったく同じというわけではありません。ただし「文脈が伸びるほど毎ステップ読み出すデータ量が増える」という構造そのものは変わらないので、今回の実測（292→119 tok/s の低下）は、このボトルネックが Apple Silicon 上でも確かに現れることを示しています。

キャッシュを小さくするもう一つの軸 — GQA

ここまでは「文脈を伸ばすとキャッシュも速度低下も避けられない」という、いわば動かせない制約の話をしてきました。ですが、キャッシュのサイズを決める式にはもう 1 つ、動かせる変数があります。KVヘッド数です。

理論式を思い出してください。2 × 層数 × KVヘッド数 × head_dim × 文脈長 × dtypeバイト数。ここまで固定して考えていた「KVヘッド数」は、実は attention head の数（このモデルでは 32）と同じである必要はありません。今回使ったモデルは n_heads=32 に対して n_kv_heads=8 で、すでに 32 個の attention head を 8 グループに分け、各グループが 1 組の Key/Value を共有する仕組みを採用しています⁴。これを Grouped-Query Attention（GQA）と呼びます。

素朴に全 head が個別の Key/Value を持つ場合（n_kv_heads=32）に比べ、このモデルの KV キャッシュはすでに 1/4 に削減されています。文脈を短くする代わりに、1 トークンあたりのキャッシュ量そのものを減らす、というアプローチです。

「文脈に対して線形に増え続けるキャッシュ」と「そのキャッシュを減らす GQA という軸」。この 2 つは、最終回で見る PHOTON の「なぜ最大 475 倍が成立するのか」を自分の言葉で分解するときの材料になります。ここでは材料を置くだけにして、分解は最終回にまとめて行います。

手元に残るもの と 次回予告

今回手を動かして残ったものを確認します。

• 動くコード — kv_cache_demo.py（mlx_lm の内部キャッシュを introspect する --probe、文脈 1K〜32K での三点一致を測る --kv、decode 速度を測る --decode）
• 測った数字 — KV キャッシュは文脈に線形（1,024→32,768 で正確に 32 倍）、256 境界での「確保容量」の 2 倍ジャンプの罠、decode 速度は 32 倍の文脈で約 2.45 倍低下（292→119 tok/s）
• メンタルモデル — KV キャッシュは attention の再計算を避ける代わりに、文脈に比例するメモリと、そのメモリを毎ステップ読み出すコストを引き受けている

次回 Part 4 では、視点を変えます。「1 回の生成に時間がかかるなら、複数回生成して多数決を取れば賢くなれるのではないか」というサンプリングの工夫（self-consistency）を実測します。今回見た「キャッシュも decode コストも避けられない」という制約を裏側に置きながら、生成の やり方 そのものを工夫する回です。

コードはサンプルコード（companion repo）の part-03 タグで再現できます（kv_cache_demo.py と、mlx 不要で通る --selftest）。clone と各 Part のタグは、companion repo かこの記事末尾のシリーズナビからたどれます。

参考文献