AIガイド · 19分で読めます

KVキャッシュとは?LLMがAttentionの状態を再利用する仕組み

KVキャッシュとは何か、LLMのDecodeを高速化する仕組み、コンテキスト長に伴うメモリの増加、プロンプトキャッシュやAPIコストとの関係をわかりやすく解説します。PrefillとDecodeの違いや、開発チームが実践できる運用上のポイントも紹介します。

色付きのトークンタイルの列が積み重なったキャッシュの板に入り、より速く光る出力ストリームが出てくるイラスト

**KV cache(Key-Valueキャッシュ)**は、大規模言語モデルが処理済みトークンについて計算した各層のKeyとValueのAttentionテンソルを保存する、推論の最適化手法です。自己回帰型のデコードでは、モデルはそれまでのすべてのトークンを各ステップで再計算する代わりに、保存した状態を再利用します。これにより、シーケンスが長くなる際の重複計算を減らせます。

ホスト型APIを呼び出す場合も、オープンウェイトモデルを自前のGPUで実行する場合も、KV cacheはレイテンシ、メモリの逼迫、スループットの中心にあります。その仕組みを理解すると、最初のトークンまでの時間とストリーミング速度の体感が異なる理由、長いコンテキストが大量のVRAMを消費する理由、APIリクエスト間のプロンプトキャッシュに完全に一致するプレフィックスが必要な理由がわかります。

KV cacheなしではAttentionはどう動く?

KV cacheがない場合、自己回帰型Transformerは新しいAIモデルのトークンを1つ生成するたびに、シーケンス全体のAttentionの射影を最初から再計算しなければなりません。

Transformerの各デコーダー層では、Self-Attentionが各入力トークンの表現を3つのベクトルに射影します。

  • Query ($Q$): 関連するコンテキストを探す現在の位置を表します。
  • Key ($K$): Queryとの照合に向けて各位置が提供するものを表します。
  • Value ($V$): QueryとKeyの内積からAttentionの重みを計算した後に集約される情報を保持します。

デコーダーのみの言語モデルはテキストを1トークンずつ生成するため、トークン$t+1$の生成には、最新トークンのQueryベクトルがトークン$1$から$t$までのKeyとValueを参照するだけで十分です。ランタイムがトークン$t$の出力後にそれらのKeyとValueを破棄すると、ステップ$t+1$では、まったく同じ$K$と$V$のテンソルを再現するだけのために、プレフィックスの各トークンをすべての層に再び通さなければなりません。プロンプトと出力が長くなるにつれ、この反復作業は生成処理全体で二次的に増大します。

なぜ因果的AttentionによってKVキャッシュが可能になる?

因果マスクにより、トークンのKeyとValueのテンソルは、それより前のトークンにのみ依存し、後から生成されるトークンには決して依存しません。

因果的なデコーダー(左から右へ処理)では、トークン$i$は位置$j > i$を参照できません。情報はシーケンス内を前方にのみ流れるため、次のことが成り立ちます。

  1. 層$l$におけるトークン$i$の表現は、トークン$1 \dots i$の処理が終わると完全に確定します。
  2. シーケンスの末尾にトークン$i+1$を追加しても、どの層でも位置$1 \dots i$の隠れ状態、Key、Valueは変化しません。
  3. ランタイムは新たに計算した$K_{i+1}$と$V_{i+1}$のベクトルを層ごとのバッファー、すなわちKV cacheに追加し、次のデコードステップには最新の1トークンだけを渡せます。

Hugging Face TransformersのKV cacheガイドとTransformersのキャッシュアーキテクチャ概要で説明されているように、キャッシュによって各デコードステップは、シーケンス全体の処理から、過去のKeyとValueをメモリから読み取る1トークンの逐次更新に変わります。

PrefillとDecodeでKV cacheの使い方はどう異なる?

LLMの推論は、PrefillとDecodeという2つの異なるフェーズに分かれます。それぞれKV cacheとのやり取りが根本的に異なり、異なるハードウェアのボトルネックに直面します。

推論フェーズ KV cacheで起きること 主なハードウェアのボトルネック ユーザーに見える指標
Prefill(プロンプト処理) すべてのプロンプトトークンの$K$と$V$を並列に計算し、各層の初期キャッシュを構築します。 計算律速(プロンプト全体にわたる大規模な行列同士の乗算)。 最初のトークンまでの時間(TTFT)。
Decode(トークン生成) 新しいトークンを1つずつ$Q$、$K$、$V$に計算し、新しい$K$/$V$のエントリーを追加して、キャッシュされた履歴全体を読み取ります。 メモリ帯域幅律速(各ステップでモデルの重みと増大するKV状態を読み込み)。 トークン間レイテンシ / 出力ストリーミング速度。

NVIDIAによるPrefill、Decode、KVキャッシュの技術的な説明のとおり、Prefillフェーズはすべての入力トークンを同時に処理します。プロンプト全体があらかじめわかっているため、GPUは射影を密な行列演算にまとめ、演算ユニットを十分に活用し、1回の処理でKV cacheを構築できます。検索拡張生成(RAG)パイプラインにおける複数文書のコンテキストなど、長いプロンプトはPrefillの作業量を直接増やし、最初のトークンまでの時間を長くします。

最初の出力トークンが出ると、推論はDecodeフェーズに入ります。各ステップではシーケンスごとに新しいトークンを1つだけ処理し、そのKeyとValueのベクトルをキャッシュに追加して、Attentionを計算するために、以前保存したすべての$K$と$V$のテンソルをGPUの高帯域幅メモリ(HBM)から読み取ります。各ステップの演算強度は低い一方、メモリの読み取り量はシーケンス長とバッチサイズに応じて増えるため、デコードは純粋な演算FLOPsよりもメモリ帯域幅と利用可能なVRAMにすぐ制約されます。

KV cacheはどれくらいのメモリを使う?

KV cacheのメモリ使用量は、シーケンス長、バッチサイズ、層数、Key-ValueのAttentionヘッド数に比例して増加します。

標準的なTransformerキャッシュで使われるテンソルの形状では、各層がバッチ、KVヘッド、シーケンス長、ヘッドの次元に沿ったKeyテンソルとValueテンソルをそれぞれ1つ保存します。これに基づき、シーケンスごとのキャッシュサイズは次のように表せます。

bytes ≈ 2 × layers × kv_heads × head_dim × seq_len × bytes_per_element

この式の各項は、モデル設計者や推論エンジニアがメモリを抑えるために調整できる要素です。

  • 2: Key(K)とValue(V)の両方のテンソルの保存を表します。
  • layers: モデル内のTransformerデコーダーブロックの総数です。
  • kv_headsとhead_dim: 各層のKey/ValueのAttentionヘッド数と各ヘッドの次元です。
  • seq_len: キャッシュされたトークンの総数です(プロンプトトークンとそれまでに生成されたトークン)。
  • bytes_per_element: 保存されたテンソルの数値精度です(例えばFP16/BF16では2バイト、8ビット形式では1バイト、NVFP4のような4ビットKV cache量子化ではさらに小さい端数のサイズ)。

kv_headsを減らす:MHA、MQA、GQA

従来のMulti-Head Attention(MHA)では、各Queryヘッドが専用のKeyヘッドとValueヘッドを持ちます(kv_heads = query_heads)。コンテキストウィンドウとバッチサイズが大きくなると、ヘッドごとに別々のKとVの射影を保存することが大きなメモリのボトルネックになりました。

次の2つのアーキテクチャ変更は、上の式のkv_headsを直接減らします。

GQAに加えて、推論エンジンはキャッシュされたKとVのテンソルを低精度の形式に量子化してbytes_per_elementも減らし、より長いコンテキストやより大きな同時バッチをGPUメモリに収めます。

サービングエンジンは大規模なKV cacheメモリをどう管理する?

本番環境のサービングエンジンは、連続したテンソルを固定サイズのブロックに分割してKV cacheメモリを管理します。これにより、GPUはVRAMを無駄にせず、より多くの同時リクエストを処理できます。

単純な実装では、サービングランタイムは各リクエストの最大可能シーケンス長に合わせて、連続したGPUメモリ領域を事前確保します。実際の出力長は事前にわからないため、この事前確保は深刻な内部・外部メモリ断片化を引き起こします。確保済みの領域が使われない一方で、新しいリクエストの実行が阻まれます。

Kwonらは、vLLMの中核となる仕組みであるPagedAttentionをEfficient Memory Management for Large Language Model Serving with PagedAttention(arXiv:2309.06180)で発表し、このボトルネックに対処しました。オペレーティングシステムの仮想メモリとページングに着想を得ています。

  • 非連続の固定サイズブロック: PagedAttentionは各リクエストのKV cacheを固定サイズのブロックに分割します。各ブロックは一定数のトークンのKeyとValueを保持し、論理ブロックから物理ブロックへの対応表で割り当てられます。
  • オンデマンド割り当て: 新しい物理ブロックは新しいトークンが生成されたときにのみ割り当てられ、内部断片化をほぼ解消します。
  • ブロック単位のプレフィックス共有: 複数のリクエストが同一のプロンプトプレフィックス(システムプロンプトや共通のfew-shot例など)を共有する場合、それぞれの論理ブロック表はメモリ内の同じ物理KVブロックを参照でき、シーケンスが分岐したときだけコピーオンライトを使います。

vLLM論文の評価(arXiv:2309.06180)では、このページ単位のKV管理により、同じGPUメモリに大幅に大きなバッチを収めることで、同程度のレイテンシでFasterTransformerとOrcaの2–4×のスループットを達成しました。vLLMやNVIDIA TensorRT-LLMのKV cache再利用システムを含む現代のサービングスタックは、ブロック単位のKV管理を利用して共通プレフィックスを再利用し、同時実行数を最大化します。

KV cacheとプロンプトキャッシュの違いは?

KV cacheは、デコードステップ間でAttentionの状態を再利用する、リクエスト単位の基盤となる推論機構です。一方、プロンプトキャッシュは、同一のプレフィックスを共有する別々の呼び出しの間でKVブロックを保持・再利用する、リクエスト横断のサービングおよびAPI機能です。

開発者はレイテンシと請求額を最適化する際に両方の用語に出会うことが多いため、1回の生成呼び出し内で起きることと、複数のAPIリクエスト間で起きることを区別すると役立ちます。

項目 KV cache(リクエスト内のデコード) プロンプトキャッシュ(リクエスト間のプレフィックス再利用)
範囲 1つのリクエスト内のトークン生成ステップ間。 短い保持期間内の複数のリクエスト間。
回避すること 新しいトークンをデコードするたびに、同じ応答内の以前のトークンを再計算すること。 別々の呼び出しで繰り返される静的なプロンプトプレフィックスに対してPrefillを再実行すること。
ライフサイクル Prefill中に作成、Decode中に追加され、リクエスト完了時に解放または退避されます。 将来の呼び出しに備え、退避ポリシーとTTLを伴ってブロックストレージ(GPUまたはホストメモリ)に保持されます。
請求上の見え方 標準の推論に組み込まれ、別項目として請求されません。 APIプロバイダーがキャッシュ済み入力トークンの割引料金や短いTTFTとして提供することが多くあります。

因果的Attentionにより、プロンプトキャッシュには厳密なルールがあります。プレフィックスは最初のトークンから完全に一致しなければなりません。各トークンのKeyとValueのテンソルは、それ以前のトークンの正確な並びに依存するため、プロンプトの先頭にあるタイムスタンプ、ユーザーID、指示のいずれか1つを変えただけでも、それ以降のすべてのトークンについてキャッシュされた$K$と$V$の状態が無効になります。

プロバイダーごとの保持ルール、最小トークン数のしきい値、割引体系については、プロンプトキャッシュの仕組みと各社の料金を解説したガイドをご覧ください。

開発者とチームに役立つ実践ポイント

KV cacheの振る舞いを理解することは、プロンプト設計の改善、レイテンシの安定化、インフラ費用の削減に直結します。

  1. 静的な内容をプロンプトの先頭に置く: システム指示、ツールスキーマ、共通の参照文書を、動的なユーザー入力やリクエスト固有のメタデータより前に配置します。同一のプレフィックスがあれば、ホスト型APIもvLLMやTensorRT-LLMのようなローカルエンジンも、重複するPrefillを実行せずにキャッシュ済みKVブロックを再利用できます。
  2. 自前で運用するGPUでコンテキスト長が同時実行数に与える影響を見る: KV cacheのサイズはseq_lenに比例して増えるため、平均コンテキスト長が2倍になると、GPUメモリが同時に保持できるリクエスト数は半分になります。自前運用向けのオープンウェイトモデルを選ぶ際は、Grouped-Query Attention(GQA)を使うアーキテクチャと、ページ単位のKV割り当ておよびKV量子化をサポートするランタイムを優先しましょう。
  3. コンテキストウィンドウが大きくても不要なコンテキストを削る: 毎ターン大量の履歴を入力すると、PrefillのレイテンシとKVメモリの圧迫が増え、トークン料金も上がります。プレフィックスキャッシュと、対象を絞った検索や定期的な要約を組み合わせましょう。これらの基本的な手法は、AI APIコストを削減する方法のガイドとAI API料金比較で紹介しています。

よくある質問

KV cacheはモデルの出力品質を変える?

モデル本来の推論精度(FP16やBF16など)で保持される標準的なKV cacheは、各ステップでプレフィックスを再計算することと数学的に等価です。因果マスクにより、将来のトークンが過去のKeyとValueのベクトルを変えることがないためです。品質とのトレードオフが生じるのは、損失のあるKV cache量子化(8ビットや4ビットの圧縮など)を適用する場合や、メモリ節約のために古いトークンを退避する場合だけです。

冒頭が変わった場合、KV cacheはなぜプロンプト途中のトークンを再利用できない?

因果的なTransformerでは、各トークンの表現に、シーケンス内のそれ以前のすべてのトークンに対するAttentionが組み込まれます。プロンプトの最初の10トークンが変わると、11番目以降のトークンの隠れ状態、ひいてはKeyとValueの射影も変わります。これらの状態は直前までの正確なプレフィックスに依存するため、サービングエンジンがリクエスト間でキャッシュ済みKVブロックを再利用できるのは、プレフィックスが最初から一致する場合だけです。

KV cacheはモデルの重みを保持する?

いいえ。モデルの重みは学習中に得られた静的なパラメーターであり、GPUが処理するすべてのリクエストで共有されます。KV cacheは、リクエスト固有の動的な活性化テンソル(現在コンテキストにあるトークンのKeyとValueの射影)を保持します。推論時のGPUメモリ全体には、静的なモデルの重み、一時的なランタイムの活性化、およびバッチ内のすべての有効なリクエストの動的なKV cacheを収める必要があります。

同時実行数が多いと、出力トークンをさらに生成する速度が遅くなるのはなぜ?

Decodeフェーズでは、新しいトークンを生成する各ステップで、バッチ内の各有効なリクエストについて、それ以前のすべてのトークンの保存済みKeyとValueのテンソルを読み取る必要があります。シーケンス長と同時リクエスト数が増えると、各ステップでGPUメモリから転送されるKVデータ量が増え、メモリ帯域幅が飽和します。その結果、スケジューラーはバッチサイズを制限するか、メモリの読み取りを待つ必要があります。

出典

執筆

Chunyu Cai

Chunyu CaiはLotook, LLCの最高経営責任者(CEO)であり、起業家です。AIモデルやエージェント、そしてそれらが日々の仕事をどう変えるかについて執筆しています。Chunyu CaiはImperial College Londonで修士号を取得しています。