AI・機械学習
上級

KVキャッシュ最適化(ケーブイキャッシュサイテキカ)

Transformerモデルの推論時にKey-Valueペアを再利用することでトークン生成速度を大幅に向上させる技術。メモリ効率と推論レイテンシのトレードオフを最適化する。

0 回閲覧
0 いいね

KVキャッシュ最適化とは

KVキャッシュ最適化(KV Cache Optimization)は、大規模言語モデル(LLM)の推論フェーズにおいて、Transformer の Self-Attention 計算で生成される Key(K)と Value(V)のテンソルをキャッシュし再利用する技術の総称です。自己回帰的なトークン生成では、新しいトークンを1つ生成するたびに過去の全トークンに対する Attention 計算が必要ですが、KVキャッシュにより過去の計算結果を保持することで冗長な再計算を排除します。

なぜKVキャッシュが重要なのか

LLM の推論コストの大部分は Attention 計算が占めます。例えば、4096トークンのコンテキストで次のトークンを生成する場合、KVキャッシュなしでは毎回4096トークン分の K・V を再計算する必要があります。KVキャッシュを使えば、新規トークンの K・V のみを計算し、過去分はキャッシュから読み出すだけで済みます。

項目KVキャッシュなしKVキャッシュあり
計算量O(n²) per tokenO(n) per token
メモリ使用量低い高い(キャッシュ分)
レイテンシ非常に高い低い
スループット低い高い

主要な最適化手法

Multi-Query Attention(MQA)

複数の Query ヘッドに対して K・V ヘッドを1つだけ共有する方式です。KVキャッシュのメモリ使用量を大幅に削減でき、Google の PaLM や Falcon が採用しています。

Grouped-Query Attention(GQA)

MQA と従来の Multi-Head Attention の中間に位置する手法です。K・V ヘッドを Query ヘッドのグループ数だけ用意します。Llama 2 70B や Mistral が採用しており、品質とメモリ効率のバランスが優れています。

KVキャッシュ量子化

KVキャッシュを FP16 から INT8 や INT4 に量子化することで、メモリフットプリントを半分以下に圧縮します。精度低下は軽微で、バッチサイズの拡大やコンテキスト長の延長が可能になります。

メモリ消費の計算例

70Bパラメータモデル(80層、64ヘッド、ヘッド次元128)で4096トークンのKVキャッシュに必要なメモリ量:

  • FP16: 80 × 2 × 64 × 128 × 4096 × 2 bytes ≈ 5.2 GB
  • GQA(8グループ): 80 × 2 × 8 × 128 × 4096 × 2 bytes ≈ 0.65 GB
  • INT8 + GQA: ≈ 0.33 GB

PagedAttention との連携

vLLM が導入した PagedAttention は、KVキャッシュを固定サイズのブロック(ページ)に分割して管理します。これにより、メモリの断片化を防ぎ、複数リクエスト間でのKVキャッシュ共有(prefix caching)も実現できます。

FAQ

Q1: KVキャッシュのサイズはモデルのどのパラメータで決まりますか?

A1: 層の数(num_layers)、K・Vヘッド数(num_kv_heads)、ヘッド次元(head_dim)、シーケンス長(seq_len)、およびデータ型のバイト数で決まります。計算式は num_layers × 2 × num_kv_heads × head_dim × seq_len × bytes_per_element です。

Q2: KVキャッシュ最適化なしでLLMは実用的に動作しますか?

A2: 短い入力(数十トークン)なら動作しますが、数千トークン以上のコンテキストでは計算量が二次関数的に増大し、実用的なレイテンシを達成できません。商用サービスではKVキャッシュは必須技術です。

Q3: Sliding Window Attentionとの関係は?

A3: Sliding Window Attention(Mistral等が採用)は、KVキャッシュの保持範囲を直近のウィンドウサイズに限定する手法です。メモリ使用量を一定に抑えつつ、長いコンテキストを処理できます。KVキャッシュ最適化の一変種と位置付けられます。