JA EN
体系推論・高速化
·無料·10分で読めます

KVキャッシュを1から理解する — 推論高速化の核心

LLMは1トークンずつ答えを吐きます。素朴に実装すると、1トークン出すたびに文全体を計算し直すという壮大な無駄が生まれる。KとVは過去のぶんが二度と変わらない——この一点だけで計算量が桁で落ちます。代わりに払うのがメモリで、その量は自分で見積もれます。バッチサイズと文脈長がなぜメモリで頭打ちになるのかまで。

対象textタスクinference

生成は、1トークンずつしか進まない

言語モデルが文章を書くとき、一気に全文が出てくるわけではありません。次に来る1トークンを予測し、それを入力の末尾に足し、また次の1トークンを予測する。この繰り返しです。自己回帰(autoregressive)生成と呼ばれる、ひどく逐次的な手続きです。

つまり100トークンの返答を作るには、モデルの順伝播を100回走らせることになります。ここに、素朴に書くと必ず踏む落とし穴があります。

素朴な実装は、毎回すべてを計算し直す

いちばん素直な実装はこうです。プロンプトを入れて1トークン得る。得たトークンをプロンプトの末尾に連結し、その全体をもう一度モデルに入れて次のトークンを得る。また連結して、また全体を入れる。

正しく動きます。動きますが、tt 番目のトークンを出すために長さ tt の系列を丸ごと処理しているので、Attentionの記事で見たとおり注意の計算だけで O(t2)O(t^2) かかります。これを t=1t=1 から nn まで足し上げると、生成全体で O(n3)O(n^3) のオーダーです。文脈が長くなるほど、この無駄は加速度的に効いてきます。

FIG 1素朴な再計算は文長の3乗、キャッシュを使うと2乗。オーダーが1つ違うだけで、長い文脈では現実的な待ち時間と非現実的な待ち時間の差になる

KとVは、過去のぶんが二度と変わらない

無駄の正体を突き止めます。自己注意は各トークンから Query・Key・Value の3つを作り、QueryとKeyの内積で重みを決めてValueを混ぜるのでした。ここで因果マスク——生成モデルは未来のトークンを見てはいけないので、各トークンは自分より前しか参照しない——という制約を思い出してください。

すると、こういうことになります。3番目のトークンの K と V は、1番目・2番目のトークンから計算されるものには一切影響しません。それどころか、4番目、5番目、100番目のトークンが後から追加されても、3番目の K と V は1ビットも変わりません。 各位置の K と V は、その位置とそれ以前の入力だけで決まる値だからです。

変わらない値を毎回計算し直しているのなら、答えは1つです。取っておいて、使い回す。 これがKVキャッシュです。

新しいトークンを1個生成するとき、必要な計算はこうなります。

1ステップあたりの注意の計算は O(t)O(t) に落ち、生成全体では O(n2)O(n^2)O(n3)O(n^3) から丸々1つオーダーが下がりました。

def decode_step(x_new, cache, layer):
    q = x_new @ layer.Wq                     # 新しい1トークンぶんだけ
    k = x_new @ layer.Wk
    v = x_new @ layer.Wv
    cache.k = concat(cache.k, k)             # 過去は再計算せず追記のみ
    cache.v = concat(cache.v, v)
    w = softmax(q @ cache.k.T / sqrt(d_k))   # 全履歴との内積
    return w @ cache.v

Q だけがキャッシュされないのは、Q は「いま何を訊きたいか」であって、毎ステップ新しいトークンのぶんしか要らないからです。

prefill と decode という2つの顔

キャッシュを入れた瞬間、推論は性格の違う2つのフェーズに分かれます。

prefill(プロンプト処理)は、入力されたプロンプト全体を1回で流し、全位置の K・V をまとめて作ってキャッシュを埋める段階です。全トークンを並列に処理できるので、GPUの演算器が埋まります。ユーザーが体感する「最初の1文字が出るまでの時間」はここで決まり、プロンプトが長いほど伸びます。

decode(逐次生成)は、1トークンずつ進む段階です。1ステップで扱う新しいトークンは1個だけなので、行列積は「行列×行列」ではなく「ベクトル×行列」になります。演算の量に対して読み込むデータの量が不釣り合いに多い——重みとKVキャッシュを毎ステップ丸ごとメモリから読む必要があるからです。結果として decode は計算律速ではなくメモリ帯域律速になります。

この非対称性を知っていると、「GPUの利用率が低いのに生成が遅い」という現象に驚かなくて済みます。演算器は暇でも、メモリバスは飽和しているのです。実務では前者を TTFT(Time To First Token)、後者を1トークンあたりの生成時間として別々に測ります。片方だけ見て「遅い」と言っても、打つべき手は決まりません。

もう1つ、地味に効いてくる帰結があります。KVキャッシュはリクエストごとの状態です。会話が続いている間、そのGPUのメモリ上に固有のキャッシュが居座る。ステートレスなWebサーバーの感覚でリクエストを別のノードへ振ると、キャッシュが無い側では prefill をやり直すことになります。LLMサービングが素朴なロードバランスと相性が悪いのは、この状態を持ってしまうためです。

キャッシュの大きさを、自分で見積もる

速くなった代わりに、メモリを食います。どれだけ食うかは掛け算で出せます。

M=2×L×n×Hkv×dh×b×BM = 2 \times L \times n \times H_{kv} \times d_h \times b \times B
(1)

これは理論というより勘定書きです。つまり「層の数だけ、いま文脈に入っているトークンの数だけ、K と V を2本ずつ抱える。それを同時に走っているリクエストの数だけ並べる」——持っているものを数え上げて掛け算にしただけのものです。

記号を1つずつ。先頭の 22 は K と V の2本ぶん。LL は層の数(各層が独自の K・V を持つ)。nn は文脈長(キャッシュ済みトークン数)。HkvH_{kv} は KV ヘッドの数、dhd_h は1ヘッドの次元で、Hkv×dhH_{kv} \times d_h が1層あたりの K のサイズです。bb は1要素のバイト数(fp16なら2)。BB は同時に処理するリクエスト数、つまりバッチサイズ。

この式が言っているのは要するに、キャッシュは文脈長にもバッチサイズにもまっすぐ比例して増えるということです。

数字を入れてみます。層32、Hkv×dh=4096H_{kv} \times d_h = 4096、fp16、文脈4096トークン、バッチ1なら

2×32×4096×4096×2 バイト=2 GiB2 \times 32 \times 4096 \times 4096 \times 2\ \text{バイト} = 2\ \text{GiB}

左辺は数式というより買い物リストです。つまり、32層それぞれで、4096トークンぶん、1トークンにつき4096個の数を、K と V の2本、1個2バイトで持つ——それを全部数え上げると右辺になります。賢いことに使われているメモリは1バイトもなく、ただ「持たなければならないもの」を持っているだけです。

1リクエストで2 GiB。バッチを8にすれば16 GiB、文脈を倍にすれば32 GiB。重みを載せた残りのメモリを、この式が食い尽くしていきます。 同時実行数と文脈長の上限がメモリで決まる、というのはこの計算の帰結にすぎません。

減らす手はいくつかあります。よく使われるのが GQA(Grouped-Query Attention)で、Query のヘッドは多いまま K・V のヘッドだけを減らして共有する設計です。式の HkvH_{kv} が直接小さくなるので、キャッシュも比例して縮みます。もう1つが KV キャッシュ自体を低精度で持つこと(bb を減らす)で、これは量子化の記事の話がそのまま当てはまります。

断片化という、もう1つの敵

素朴なサーバーは、リクエストごとに「最大文脈長ぶん」の連続領域を先に確保します。実際には200トークンで終わる返答にも4096トークンぶんの部屋を取ってしまうので、大半が使われないまま塞がれる。OSの仮想メモリと同じ発想で、キャッシュを固定長のブロックに切って必要なぶんだけ割り当てる方式(PagedAttention として知られ、vLLM が実装しています)が広く使われるようになったのはこのためです。同じプロンプト前半を共有する複数リクエストで、ブロックを共有できるという副産物もあります。

現場ではこう使う

KVキャッシュが仕事に直結するのは、LLMを自前でサービングするとき——GPUを何枚買うか決める人、同時接続数のSLOを引く人、レイテンシの苦情を調べる人です。API を叩くだけの立場でも、料金と速度の振る舞いを説明するのに必要になります。

触るパラメータ。 vLLM なら --max-model-len(文脈長の上限)、--max-num-seqs(同時実行数)、--gpu-memory-utilization(重み+KVに使ってよいメモリ比率)、--kv-cache-dtype(KVの精度)、--enable-prefix-caching(共通プレフィックスのKVを使い回す)。Hugging Face transformers なら use_cachepast_key_values が同じ概念の生の姿です。

いちばん多い事故が「起動したのにOOMで落ちる」。 重みが載ったことと、本番の同時実行数×文脈長でKVが載ることは別問題です。空いているメモリを 2LHkvdhb2 L H_{kv} d_h b で割れば「同時に抱えられる総トークン数」が出ます。この1本の割り算を先にやっておくかどうかで、深夜の障害対応が1回減ります。

プレフィックスキャッシュを自分で壊さない。 システムプロンプトの先頭に現在時刻やリクエストIDを入れると、リクエストごとに全プレフィックスが別物になり、共有できるはずのKVが毎回作り直されます。可変部分は後ろに置く——これだけで prefill の負荷が変わります。

遅さの原因を、正しい側で探す。 最初のトークンまでが遅いなら prefill、つまりプロンプト長の問題です。出力がだらだら遅いなら decode、つまりメモリ帯域とKVの読み出しの問題で、プロンプトを短くしても直りません。前者は入力を削る・プレフィックスを共有する、後者はバッチングやKVの精度を見直す、と打ち手が変わります。

面接や設計レビューで問われるのは「文脈長を2倍にすると何が2倍になるか」です。KVキャッシュのメモリは線形に2倍、しかし prefill の注意計算は2乗で4倍。この2つを混同している回答は多く、区別できるだけで理解の深さが伝わります。

まとめ

次は、そのメモリをさらに削る側の話——重みとキャッシュを低精度で持つ量子化です。

コメント

コメントにはログインが必要です