@FukasawaYusuke
KV Caching Explained: Optimizing Transformer Inference Efficiency

Prefix caching | LLM Inference Handbook
前提: Transformer の Self-Attention では、各トークンについて Q(Query)、K(Key)、V(Value) の3つのベクトルを計算する
Attention(Q, K, V) = softmax(QK^T / √d_k) V
KV Cache
K と V の値は、そのトークンより前のトークンにのみ依存する(decoder modelなので時間方向は一定)。つまり、一度計算した K・V は後から変わらない
なので、以下のようにすれば再計算を防げる
ステップ1: トークン [A, B, C] → K₁,K₂,K₃ と V₁,V₂,V₃ を計算 → GPU メモリに保存 ステップ2: トークン D を生成 → K₄,V₄ だけ新規計算、K₁~K₃,V₁~V₃ はキャッシュから再利用 ステップ3: トークン E を生成 → K₅,V₅ だけ新規計算、K₁~K₄,V₁~V₄ はキャッシュから再利用
これにより、各ステップの計算量が O(n²) → O(n) に削減!
Prefix Caching
KV Cache は通常「1つのリクエスト内」で使われるが、Prefix Caching はこれをリクエスト間に拡張したもの
たとえば、以下のような3つのリクエストで「システムプロンプト」部分は同一な状況を考える
リクエスト1: [システムプロンプト] + [ユーザーA の質問] リクエスト2: [システムプロンプト] + [ユーザーB の質問] リクエスト3: [システムプロンプト] + [ユーザーC の質問]
この部分の KV Cache を一度計算して保存しておけば:
リクエスト1: システムプロンプトの KV を計算 → 保存 → ユーザーA 分だけ追加計算 リクエスト2: 保存済み KV を再利用 → ユーザーB 分だけ追加計算 ← prefill スキップ リクエスト3: 保存済み KV を再利用 → ユーザーC 分だけ追加計算 ← prefill スキップ
prefill(入力プロンプト全体の K・V を計算するフェーズ)をスキップできる!
なぜ「プレフィックス(先頭)」限定?
LLM 系の Self-Attention の caual mask の性質により、位置 i のトークンの KV は、位置 0~(i-1) のトークンにのみ依存する
位置: 0 1 2 3 4 5 6 [sys] [sys] [sys] [sys] [user] [user] [user] ↓ ↓ ↓ ↓ K₀は独立 K₁はK₀に依存 K₂はK₀,K₁に依存 ...
もし途中のトークンが変わると、それ以降の全 KV が無効
だから先頭から連続して一致している部分だけがキャッシュとして再利用可能
@MizutaYuki
「定時実行」と「定期実行」の実装ガイド | blog.jxck.io