AI・機械学習
上級

KVキャッシュ管理(ケーブイキャッシュカンリ)

KVキャッシュ管理は、Transformer系LLMの推論時に各Attention層で生成されるKey-Valueペアを効率的に格納・再利用・共有するメモリ管理技術。PagedAttention・Prefix Caching・Radix Attention・Token Eviction(トークン退避)などの手法があり、GPUメモリの利用効率を60〜80%改善して同時処理可能なリクエスト数を数倍に増加させる。LLM推論エンジンのスループットとコスト効率を決定する最重要コンポーネント。

0 回閲覧
0 いいね

KVキャッシュとは

KVキャッシュ(Key-Value Cache)は、TransformerアーキテクチャのSelf-Attention計算において、過去のトークンに対するKey(K)ベクトルとValue(V)ベクトルを保存・再利用する仕組みである。自己回帰(Autoregressive)生成では、新しいトークンを1つ生成するたびに過去の全トークンとのAttentionスコアを計算する必要がある。KVキャッシュがなければ、n番目のトークン生成時に1番目からn-1番目のトークンのK/Vを再計算しなければならず、計算量はO(n²)に爆発する。KVキャッシュは過去のK/Vを保存して再利用することで、各ステップの計算量をO(n)に削減する。

KVキャッシュのメモリ消費量

KVキャッシュのサイズは「層数 × 2(K+V)× ヘッド数 × ヘッド次元 × シーケンス長 × バッチサイズ × データ型サイズ」で決まる。具体例として、Llama 3.1 70B(80層、64ヘッド、ヘッド次元128)でFP16(2バイト)、シーケンス長8192の場合:

80 × 2 × 64 × 128 × 8192 × 2 = 約21.5 GB / リクエスト

1リクエストだけで21.5GBのKVキャッシュが必要であり、A100 80GBではモデル重み(約140GB、TP=2で70GB/GPU)を差し引くと、1GPU あたり1〜2リクエストしか同時処理できない計算になる。KVキャッシュ管理の効率化は、LLM推論のスケーラビリティを直接左右する最重要課題である。

モデルパラメータ数層数ヘッド数KVキャッシュ/リクエスト (8K, FP16)同時処理数 (A100 80GB)
Llama 3.1 8B8B3232約2.1 GB8〜16
Llama 3.1 70B (TP=4)70B8064約21.5 GB1〜2
Mixtral 8×7B47B3232約2.1 GB (GQA)8〜16
GPT-4クラス1.8T推定120+128+約100 GB+<1 (TP=8必須)

PagedAttention

PagedAttentionは、vLLMで導入されたKVキャッシュのページング管理技術である。OSの仮想メモリシステムから着想を得て、KVキャッシュを固定サイズの「ブロック」(デフォルト16トークン分)に分割し、非連続なGPUメモリ領域に動的に配置する。

従来方式の問題

従来のLLM推論では、各リクエストに対してモデルの最大シーケンス長分のKVキャッシュメモリを連続領域で事前確保していた。しかし実際の生成長は事前に予測できないため、以下の無駄が生じる:

  • 内部断片化: 最大8192トークン分確保したが、実際に128トークンしか生成しない → 98%のメモリが無駄
  • 外部断片化: リクエストの開始・完了が繰り返されると、メモリの空き領域が細かく分断される → 大きな連続領域を確保できない

PagedAttentionの解決策

  1. KVキャッシュを16トークン単位のブロックに分割
  2. ブロックテーブル(ページテーブル相当)で論理ブロック→物理ブロックのマッピング管理
  3. トークン生成の進行に応じてブロックを1つずつ動的確保
  4. リクエスト完了時にブロックを即座にフリーリストに返却
  5. Copy-on-Write(CoW)でBeam Search/並列サンプリング時のブロック共有
指標従来方式PagedAttention
メモリ利用効率20〜40%90〜95%
同時処理可能リクエスト数N3〜5N
メモリ確保遅延大(連続領域探索)小(ブロック単位)
Beam Search メモリ消費ビーム数×フルコピーCoWで物理共有

Prefix Caching(プレフィックスキャッシュ)

Prefix Cachingは、共通プレフィックス(システムプロンプト、Few-Shotの例示)のKVキャッシュを複数リクエスト間で物理的に共有する技術である。

ユースケース

多くのLLMアプリケーションでは、システムプロンプト(数百〜数千トークン)が全リクエストで共通である。同一システムプロンプトのKVキャッシュを1度計算すれば、以降のリクエストでは再計算せずにキャッシュを参照するだけでよい。これにより:

  • TTFT(Time to First Token)削減: 共通プレフィックスのPrefill計算をスキップ。2048トークンのシステムプロンプトで300〜500ms短縮。
  • メモリ節約: 共通プレフィックスのKVキャッシュを物理的に1コピーで共有。100リクエストで99コピー分のメモリを節約。

SGLangのRadixAttention

SGLangはRadix Tree(基数木)データ構造でKVキャッシュのプレフィックスを管理する。文字列のプレフィックス共有をトークンシーケンスに適用し、分岐点からのみ新規KVキャッシュを計算する。Few-Shot推論(同一例示セットで入力のみ変更)では、Prefix Cachingなしと比較してTTFTが5〜10倍改善する。

Automatic Prefix Caching(APC)

vLLMの--enable-prefix-cachingオプションで有効化される自動プレフィックスキャッシュ。リクエスト間でプレフィックスの一致をハッシュベースで自動検出し、明示的な設定なしにキャッシュ共有を行う。

トークン退避と長コンテキスト対応

128K〜1Mトークンの長コンテキスト推論では、KVキャッシュがGPUメモリを圧迫する。以下の技術で対応する。

KVキャッシュ圧縮

  • GQA(Grouped Query Attention): Keyヘッドを複数Queryヘッドで共有し、KVキャッシュサイズを1/4〜1/8に削減。Llama 3、Mistral、Gemmaが採用。
  • MQA(Multi-Query Attention): 全Queryヘッドで1つのKey/Valueを共有。GQAのさらなる圧縮版。Falcon、StarCoder等。
  • KVキャッシュ量子化: K/VベクトルをFP16からINT8/INT4に量子化。精度劣化は0.5〜1%程度でメモリ50〜75%削減。

CPU/ディスクオフロード

GPUメモリに収まらないKVキャッシュをCPUメモリやNVMe SSDに退避する。アクセス頻度の低い古いトークンのKVキャッシュを退避し、必要時にGPUに戻す。InfiniteTransformerやStreamingLLMの考え方を応用し、Attention Sinkトークン(先頭トークン)とスライディングウィンドウ内トークンのKVキャッシュのみGPU上に保持する。

Token Dropping / Eviction

重要度の低いトークンのKVキャッシュを破棄する手法。Attention重みの累積値が閾値以下のトークンを「重要度低」と判定し、そのK/Vペアを解放する。H2O(Heavy-Hitter Oracle)やScissorHandsなどのアルゴリズムが提案されている。精度への影響は用途依存であり、要約タスクでは影響が小さいが、コード生成や数学推論では劣化が顕著。

手法メモリ削減率精度影響実装複雑度対応エンジン
GQA75〜87.5%なし(モデル設計)モデル依存全エンジン
KV量子化 (INT8)50%0.5〜1%低vLLM, TRT-LLM
CPU Offload∞(CPU RAM分)なし(遅延増加)中vLLM
Token Dropping30〜50%1〜5%高研究段階

よくある質問(FAQ)

Q1: KVキャッシュはモデルの重みよりメモリを消費しますか?

長いコンテキストでは消費する。Llama 3.1 70B(FP16)のモデル重みは約140GBだが、128Kコンテキストの1リクエストのKVキャッシュは約340GBに達する。そのため、長コンテキスト推論ではKVキャッシュ管理がメモリのボトルネックとなり、GQA/MQAやKV量子化が不可欠である。8Kコンテキストであれば約21GB/リクエストでモデル重みより小さい。

Q2: Prefix Cachingはどんなワークロードで効果がありますか?

同一システムプロンプトまたはFew-Shot例示を使い回すワークロードで最大効果を発揮する。チャットボット(システムプロンプト共通)・バッチ評価(同一プロンプトテンプレートで入力のみ変更)・RAG(同一検索結果をコンテキストに持つ複数質問)が典型例。逆に、毎回異なるプロンプトを使う場合はキャッシュヒット率が低く、効果は限定的。

Q3: GQAとMQAはどちらが主流ですか?

2025〜2026年時点ではGQAが主流。Llama 3シリーズ、Mistral/Mixtral、Gemma 2、Qwen 2.5など主要モデルの大半がGQAを採用。MQAは圧縮率が高いが精度劣化がやや大きく、Falcon 40B等の一部モデルにとどまる。GQAはグループサイズ(通常4〜8)を調整することで圧縮率と精度のバランスを取れる柔軟性が評価されている。

Q4: KVキャッシュの量子化はモデル重みの量子化と併用できますか?

併用可能で、メモリ効率をさらに向上させる。例えば、モデル重みをAWQ 4bit(75%削減)、KVキャッシュをINT8(50%削減)とすることで、Llama 3.1 70BをA100 80GB×2構成でサービング可能にする。vLLM 0.7以降で--kv-cache-dtype fp8オプションが利用でき、FP8 KVキャッシュでメモリ50%削減・精度劣化0.3%未満を達成する。