KV Cacheとは？LLMの生成を高速化する仕組み

KV Cache（Key-Value Cache、Attentionで計算済みのKeyとValueを保存する領域）は、LLMが次のtokenを生成するたびに、過去tokenのKeyとValueを計算し直す無駄を減らす仕組みです。

生成は速くなりますが、無料ではありません。会話が長くなるほどcacheは増え、GPUメモリの容量とメモリ帯域を使います。KV Cacheは「計算を消す技術」ではなく、過去の計算結果をメモリへ保存し、再計算を読み出しへ置き換える技術です。

LLMのtoken生成とともにKeyとValueのcacheが伸びるKV Cacheのアイキャッチ

3文要約

Decoder-only LLMは、過去tokenから計算したKeyとValueを各層で保存し、次のtoken生成時に再利用します。
新しいtokenではQuery、Key、Valueを1 token分だけ計算できますが、Queryと全過去Keyのscore計算、全過去Valueの読み出しは残ります。
KV Cache容量は、概ねbatch、系列長、層数、KV head数、head dimension、保存dtypeに比例します。

代表的な論文

項目	内容
論文タイトル	Fast Transformer Decoding: One Write-Head is All You Need
著者	Noam Shazeer
公開年	2019年
論文リンク	arXiv:1911.02150
本記事との関係	逐次decodeでK/Vの読み出しがメモリ帯域負荷になることを分析し、K/Vをhead間で共有するMQAを提案

KV Cacheそのものを初めて提案した論文として扱うのではなく、逐次推論におけるK/V tensorの構造とボトルネックを確認する代表資料として参照します。

KV Cacheは何を保存するのか

Self-Attention（同じtoken列の位置同士を関連付ける処理）では、各tokenの表現からQuery、Key、Valueを作ります。

Query（Q）: 今のtokenが何を探しているか
Key（K）: 過去tokenがどの情報に対応するか
Value（V）: 注目したときに取り出す内容

自己回帰生成では、tokenを1つ追加するたびに次のtokenを予測します。過去tokenのKとVは、未来にtokenが追加されても変わりません。そのため、一度計算したKとVを各Transformer層で保存できます。

一方、過去のQは通常保存しません。次のstepで必要なのは、新しいtokenから作ったQueryと、過去を含むKey・Valueだからです。

Hugging Face TransformersのCaching解説でも、causal attention（未来tokenを参照しないAttention）では、処理済みtokenのK/Vを再利用できることが説明されています。

cacheなし・ありで何が変わるのか

LLM は文章を という4 tokenの続きとして、次のtokenを生成する場面を考えます。

KV Cacheがない場合、モデルは生成stepごとに、それまでのtoken列を先頭から通し直します。

生成step	入力済みtoken	過去K/Vの扱い
1	`LLM`	`LLM`のK/Vを計算
2	`LLM は`	`LLM`を含む2 token分を再計算
3	`LLM は文章`	過去3 token分を再計算
4	`LLM は文章を`	過去4 token分を再計算

同じtokenのK/V projection（隠れ状態をK/Vへ写す行列演算）が何度も現れます。

KV Cacheがあれば、過去tokenのK/Vは保存済みです。各stepでは新しく追加されたtokenのQ/K/Vだけを計算し、新しいK/Vをcacheの末尾へ追加します。

比較軸	KV Cacheなし	KV Cacheあり
過去tokenのK/V projection	毎step再計算	保存済み結果を再利用
新tokenのQ/K/V	計算する	計算する
過去K/VとのAttention	計算する	計算する
追加メモリ	小さい	系列長に応じて増える
主な利点	実装が単純	逐次生成の重複計算を減らす

KV Cacheなしでは過去KとVを再計算し、ありでは新token分だけ追加する比較

Attention式で見るKV Cache

時刻 \(t\) の新しいtokenから作るQueryを \(q_t\) とします。過去を含むKeyとValueは、次のように連結した形で使います。

K_{1:t}=[K_{1:t-1};k_t],\qquad V_{1:t}=[V_{1:t-1};v_t]

現在tokenのAttention出力は、簡略化すると次の式です。

o_t=\mathrm{softmax}\left(\frac{q_tK_{1:t}^{\mathrm{T}}}{\sqrt{d_{head}}}\right)V_{1:t}

cacheへ保存するのは \(K_{1:t}\) と \(V_{1:t}\) です。次のstepでは \(k_{t+1}\) と \(v_{t+1}\) だけを追加します。

K_{cache}\leftarrow\mathrm{concat}(K_{cache},k_{t+1}),\qquad V_{cache}\leftarrow\mathrm{concat}(V_{cache},v_{t+1})

この処理は1層だけではありません。各Transformer層は異なる隠れ状態からK/Vを作るため、層ごとに独立したcacheを持ちます。

Attentionの基本式はAttention Is All You Need、逐次推論でK/Vを保持するtensor形状はFast Transformer Decoding: One Write-Head is All You NeedとHugging Face公式資料で確認できます。

prefillとdecodeで役割が違う

LLM推論は、prefill（入力promptをまとめて処理する段階）とdecode（出力tokenを1つずつ生成する段階）に分けると整理しやすくなります。

段階	入力	KV Cacheの動き	並列性
prefill	prompt全体	入力token分のK/Vを各層で作る	token位置をまとめて処理しやすい
decode	新しい1 token	新token分のK/Vを末尾へ追加	token間は逐次になる

prefillでprompt分のKV Cacheを作りdecodeで1 tokenずつ追加する流れ

たとえば1,000 tokenのpromptを入力した場合、prefill後のcacheには1,000 token分のK/Vがあります。1 token生成すると1,001 token分、さらに生成すると1,002 token分へ増えます。

prefillにも計算時間は必要です。ただし入力が確定しているため、複数位置をまとめて行列計算できます。decodeは前stepの出力が次stepの入力になるので、生成token間を同じ方法では並列化できません。

前提となるDecoder-only構造は、Decoder-only Transformerとは？GPT系LLMの学習と生成を図解で詳しく説明しています。

KV Cacheがあっても生成はO(1)にならない

KV Cacheが省くのは、主に過去tokenのK/V projectionと、それらを作るために過去tokenを再び全層へ通す処理です。

一方、時刻 \(t\) のQuery \(q_t\) は、\(t\) 個のKeyとscoreを計算します。得られたattention weightを使って、\(t\) 個のValueも読み出します。contextが長くなるほど、1 token生成時に参照するcacheは増えます。

処理	cacheで省けるか	context長との関係
過去tokenのK/Vを再生成	省ける	過去分は再計算しない
新tokenのQ/K/Vを生成	省けない	1 token分は必要
\(q_tK_{1:t}^{\mathrm{T}}\)	省けない	Key数に応じて増える
weightと\(V_{1:t}\)の積	省けない	Value数に応じて増える
cacheの読み出し	省けない	cacheが長いほど増える

KV Cacheで再計算は減るがcontext長とともに容量と読み出し量が増える関係

そのため、長いcontextのdecodeでは計算能力だけでなくmemory bandwidth（単位時間にメモリから読み書きできる量）が効きます。

Fast Transformer Decoding: One Write-Head is All You Needは、incremental decoding（tokenを逐次生成する推論）で大きなK/V tensorを繰り返し読むメモリ帯域負荷を問題として示しました。KV Cacheを使えば何token目でも同じ時間になる、という主張ではありません。

KV Cacheのメモリ量を計算する

標準的なdense KV Cacheでは、KとVそれぞれの1層分のtensor shapeを次のように考えられます。

[batch size, KV head数, sequence length, head dimension]

KとVの2つがあり、それを全層で持つため、概算式は次の形です。

M_{KV}=B\times T\times L\times2\times H_{KV}\times d_{head}\times s

記号	意味
\(B\)	同時に保持するsequence数
\(T\)	promptと生成済みtokenを含む系列長
\(L\)	KV Cacheを持つAttention層数
\(2\)	KeyとValueの2 tensor
\(H_{KV}\)	Key/Value head数
\(d_{head}\)	1 headの次元数
\(s\)	1要素あたりのbyte数

具体値を代入する

特定の製品名ではなく、次の仮想的なMHA（Multi-Head Attention、各Query headが独立したK/V headを持つ構成）を考えます。

32層
32 KV heads
head dimension 128
sequence length 4,096
batch size 1
FP16またはBF16で1要素2 bytes

1\times4096\times32\times2\times32\times128\times2 =2,147,483,648\ \mathrm{bytes}

これは2 GiBです。batch sizeが8で、8 requestすべてを4,096 tokenまで保持すれば、単純合計は16 GiBになります。

32層32 KV headで4096 tokenを保持するとKV Cacheが2 GiBになる計算

この値はKV Cacheだけです。モデル重み、activation（一時的な中間値）、CUDA context、workspace、メモリアロケータの予約領域は別に必要です。したがって「24GB GPUなら、24GBすべてをKV Cacheへ使える」とは計算できません。

KV head数を減らすとcacheも小さくなる

KV Cacheの式にはQuery head数ではなく、KV head数 \(H_{KV}\) が入ります。

Attention方式	Query head	KV head	cache容量の傾向	主なトレードオフ
MHA	複数	Query headと同数	大きい	headごとに独立したK/Vを持つ
GQA	複数	Query headより少ない複数	中間	品質と速度・容量の折衷
MQA	複数	1	小さい	K/V共有が最も強い

先ほどの仮想構成で、他条件を変えずKV head数だけを変えると次の概算になります。

KV head数	4,096 token・batch 1のcache
32	2 GiB
8	512 MiB
1	64 MiB

これは容量式の比較であり、品質や実測速度がこの比率で変わる保証ではありません。

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointsは、MHAとMQAの中間としてGQAを提案し、KV head数と品質・推論速度のトレードオフを扱っています。MHA、MQA、GQAの構造はシリーズ第20回で詳しく扱います。

KV Cacheは1種類ではない

実際の推論ライブラリでは、用途に応じてcacheの確保・保存方法を選びます。

方式	概要	利点	注意点
Dynamic Cache	生成に合わせて伸ばす	無駄な事前確保を減らせる	shapeが変わり続ける
Static Cache	最大長を先に確保	固定shapeを使う最適化と相性がよい	未使用領域も確保する
Sliding Window	直近の一定範囲を保持	cache増加を上限で止められる	古いtokenを直接参照できない
Quantized Cache	K/Vを低bitで保存	容量と帯域を減らせる	量子化・復元costと誤差がある
Offloaded Cache	CPU memoryなどへ退避	GPUメモリ不足を緩和できる	転送遅延が増える

Hugging Face TransformersのCache strategiesでは、Dynamic、Static、Quantized、offloadingなどのcache方式が整理されています。対応方式はmodelやライブラリversionで変わるため、実装時は使用versionの公式資料を確認してください。

PagedAttentionはcacheの中身ではなく配置を変える

複数requestを同時処理すると、それぞれの生成長は異なります。最大長を一続きの領域として先に確保すると未使用領域が生まれ、動的に確保するとfragmentation（空き領域が細かく分かれ、使いにくくなる状態）が起こります。

Efficient Memory Management for Large Language Model Serving with PagedAttentionは、KV Cacheを固定token数のblockへ分け、物理メモリ上で連続していなくても参照できるようにしました。OSのvirtual memory（仮想記憶）とpagingを参考にした設計です。

PagedAttentionはK/Vを近似して小さくする量子化ではありません。主に配置、割り当て、共有を改善し、同じGPUメモリへより多くのrequestを収めやすくします。詳細はシリーズ第22回のvLLM・PagedAttention編で扱います。

学習時に同じcacheを使わない理由

KV Cacheは、過去tokenの状態を固定して再利用できる自己回帰推論向けの最適化です。

学習時は、正解token列をまとめて入力し、causal maskで未来を隠しながら複数位置のlossを並列に計算します。parameter更新のためにactivationと勾配も保持するので、推論時と目的が異なります。

前stepの古いparameterで作ったK/Vを次の学習stepへ持ち越せば、現在のparameterと一致しません。Hugging Face公式Caching解説も、cacheは推論だけで使い、trainingで有効にすると予期しない問題が起こり得ると注意しています。

容量式がそのまま当てはまらない場合

基本式は見積もりの出発点ですが、実際の使用量は次の条件で変わります。

beam searchで1 requestから複数候補を保持する
prefix cacheを複数requestで共有する
tensor parallelismでKV headやcacheを複数GPUへ分割する
cross-attentionを持つencoder-decoder modelを使う
sliding-window、sparse attention、hybrid modelを使う
K/Vを量子化、圧縮、CPUへoffloadする
allocatorの予約領域やblock単位の端数がある

容量計画では、式で理論値を出した後、実際のmodel configと推論engineのmetricsで確認します。

よくある誤解

KV Cacheは会話の文章そのものか

文章やtoken IDをそのまま保存する領域ではありません。各層がtoken表現から計算したK/V tensorです。同じ文章でも、modelやparameterが違えば値も変わります。

context windowとKV Cacheは同じか

同じではありません。context windowはmodelが扱えるtoken範囲の上限です。KV Cacheは、実際に処理済みのtokenについて推論中に保持する中間状態です。

cacheを使えばAttention計算は不要か

不要になるのは過去K/Vの再計算です。現在Queryと過去K/Vを使うAttention自体は必要です。

cacheを大きくすれば必ず速くなるか

長いcontextを保持できますが、読み出すK/Vも増えます。容量不足でoffloadが起きたり、memory bandwidthが詰まったりすれば遅くなる場合があります。

cacheの有無で出力は必ずbit単位で同じか

数学的には同じAttentionを効率よく計算することが目的です。ただし並列計算の形やkernelが変わると、浮動小数点の丸め順序により小さな差が出る場合があります。量子化cacheでは近似誤差も考慮します。

まとめ

KV Cacheは、Decoder-only LLMの逐次生成で、過去tokenのKeyとValueを各層に保存して再利用する仕組みです。過去K/Vのprojectionを繰り返す無駄を減らせるため、通常は生成を大きく効率化できます。

一方で、新しいQueryと全過去K/VのAttention計算、cacheの読み出し、系列長に比例するメモリは残ります。導入時は「cacheを使うか」だけでなく、KV head数、dtype、最大系列長、batch、割り当て方式まで確認します。

次回は、KV Cacheを含むLLMのメモリをさらに減らす量子化を扱います。