KVキャッシュ

KV Cacheとは何か?

KV Cache(Key-Value Cache)は、トランスフォーマーモデル、特に大規模言語モデル(LLM)における推論時の最適化手法です。これは、注意機構(Attention)の計算時に、以前に処理されたすべてのトークンに対して計算されたキー(K)とバリュー(V)のテンソルを保存します。推論の各ステップでこれらのテンソルを再計算する代わりに、モデルはキャッシュから再利用し、新しいトークンのKとVのみを計算します。このアプローチは、効率的で高速な自己回帰型テキスト生成の基盤となっています。

KV Cacheは、以前のトークンからの中間的なキーとバリューのテンソルを保持する補助メモリとして機能します。新しいトークンが生成される際、新しいトークンのKとVのみを計算して追加すればよく、以前のすべてのK/Vはキャッシュから即座に利用できるため、計算オーバーヘッドが劇的に削減されます。

KV Cacheの使用方法

KV Cacheは、トランスフォーマーベースのモデルにおいて、トークンごとにテキストを生成する推論時にのみ使用されます。

基本的な使用パターン

自己回帰型生成では、LLMが一度に1つのトークンずつテキストを生成し、各予測は以前のすべてのトークンを条件としています。各推論ステップで、モデルは次のトークンの注意を計算するために、完全なシーケンスのKとVを必要とします。KV Cacheを使用すると、すべてのステップで以前のすべてのトークンのKとVを再計算する代わりに、新しいトークンのKとVのみを計算してキャッシュに追加します。

結果: 計算量の劇的な削減、レイテンシの低減、推論時のコスト削減—特に長いシーケンスにおいて顕著です。

一般的な使用コンテキスト

• LLMによるテキスト生成(GPT、Llama、Claude、Gemini)
• チャットボットと対話エージェント
• コード補完とコードアシスタント
• リアルタイムおよびバッチ推論サービング

なぜKV CacheがLLMにとって重要なのか

KV Cacheがない場合の課題

トランスフォーマーの注意機構は、トークンごとに3つの射影を含みます:

• Query(Q): 現在のトークンが「知りたいこと」
• Key(K): 各トークンの「アドレスラベル」
• Value(V): 各トークンの「コンテンツ」

推論中、LLMは入力を一度に1つのトークンずつ処理します。標準的な推論では、すでに処理されたトークンを含む現在のシーケンス内のすべてのトークンに対してKとVを再計算します。これは長いシーケンスにとって非常に非効率的です。

非効率性の例

「The cat sits」を生成する場合:

• 「The」を生成: 「The」のK/Vを計算
• 「The cat」を生成: 「The」と「cat」の両方のK/Vを再計算
• 「The cat sits」を生成: 3つすべてのトークンのK/Vを再計算

KV Cacheを使用した場合

• 「The」のK/Vを一度計算し、保存
• 「cat」では、「cat」のK/Vを計算し、キャッシュに追加
• 「sits」では、「sits」のK/Vを計算し、キャッシュに追加;「The」と「cat」のK/Vは再利用

最適化の利点

• 速度: 最大5〜20倍高速な推論
• コスト: 計算とAPIコストの大幅な削減
• スケーラビリティ: 長いコンテキストと複数ターンの会話を可能に

KV Cacheの動作原理:詳細な例

KV Cacheなしの場合

プロンプト: ["The", "cat", "sits"]

各ステップで現在のシーケンス内のすべてのトークンのKとVを再計算:

ステップ1: "The"           --> K1, V1    (計算)
ステップ2: "The cat"       --> K1, V1, K2, V2  (K1, V1を再計算)
ステップ3: "The cat sits"  --> K1, V1, K2, V2, K3, V3 (すべて再計算)

KV Cacheありの場合

各ステップで新しいトークンのK/Vのみを計算/追加;キャッシュが以前のK/Vを保持:

ステップ1: "The"      --> K1, V1    (キャッシュに保存)
ステップ2: "cat"      --> K2, V2    (キャッシュに追加)
ステップ3: "sits"     --> K3, V3    (キャッシュに追加)

ステップ3後のキャッシュ:

K-cache: [K1, K2, K3]
V-cache: [V1, V2, V3]

技術的実装

数学的定式化

n個の入力トークンのシーケンスに対して、トランスフォーマー層は以下を計算します:

Attention(Q, K, V) = softmax(Q K^T / √d_k) V

• 訓練中: すべてのトークンのQ、K、Vを並列計算
• KV Cacheを使用した推論中: 新しいトークンのQ、K、Vのみを計算;以前のKとVはキャッシュから取得

PyTorch KV Cacheの例

class KVCache:
    def __init__(self):
        self.cache = {"key": None, "value": None}

    def update(self, key, value):
        if self.cache["key"] is None:
            self.cache["key"] = key
            self.cache["value"] = value
        else:
            self.cache["key"] = torch.cat([self.cache["key"], key], dim=1)
            self.cache["value"] = torch.cat([self.cache["value"], value], dim=1)

    def get_cache(self):
        return self.cache

Hugging Face Transformersの使用方法

最新のライブラリのほとんどは、KV Cacheを自動的に処理します:

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('gpt2')
model = AutoModelForCausalLM.from_pretrained('gpt2')

tokens = tokenizer.encode("The cat sits", return_tensors="pt")
output = model.generate(tokens, max_new_tokens=10, use_cache=True)
print(tokenizer.decode(output[0]))

use_cacheパラメータがKVキャッシングを有効にします(デフォルト: True)。

パフォーマンスへの影響

定量的ベンチマーク

ベンチマーク(T4 GPU、SmolLM2-1.7B):

• 標準推論(KV Cacheなし): 61秒
• KVキャッシング有効: 11.7秒
• 約5.2倍の高速化

APIコスト削減: 多くのプロバイダー(Anthropic、OpenAI)は、キャッシュされたトークンに対して低料金を請求—最大10倍安価(例:100万トークンあたり$0.30対$3)。

効率を最大化するためのベストプラクティス

プロンプトエンジニアリングとコンテキスト管理

安定したプロンプトプレフィックス: プロンプトプレフィックスはターン間で同一である必要があります。変更(単一トークンでも)があると、その時点からキャッシュが無効になります。

追加専用コンテキスト: 常に新しい情報を追加;以前のコンテキストを書き換えたり並べ替えたりしないでください。

決定論的順序: 構造化データは、偶発的なキャッシュ無効化を避けるために一貫した順序を持つ必要があります。

明示的なキャッシュブレークポイント: 複数ターンのエージェントの場合、コンテキストが変更される場所をマークして、フレームワークが効率を維持できるようにします。

本番環境でのキャッシュ管理

キャッシュサイズ: K/Vテンソルはコンテキスト長に比例して線形にスケールします。非常に長いシーケンスはメモリのボトルネックになる可能性があります。

キャッシュの有効期間: コンテキストが変更されたとき、またはメモリを解放する必要があるときに、キャッシュエントリを無効化/期限切れにします。

並行性: 各並行リクエストは独自のキャッシュスペースを必要とする場合があります。

ベストプラクティス: プロンプトプレフィックスを安定させ、コンテキストを追加専用に保ちます。古いエントリを変更する動的コンテキストを避けます。

高度なKV Cache技術

Paged Attention

コンテキストを「ページ」(チャンク)に分割します。関連するページのみを高速メモリに保持;古いページはオフロードまたは必要に応じて再計算されます。GPUメモリを使い果たすことなく、非常に長いコンテキスト(数万トークン)の処理を可能にします。

Radix Attention

キャッシュされたトークンをラディックスツリー構造で整理します。対数スケーリング:トークンのグループに階層的に注意を向けます。

Multi-Query Attention

注意ヘッド間でキー/バリューを共有することで、KV Cacheメモリを削減します。

新たなトレンド

• 予測的キャッシュウォーミング:予想されるニーズに基づいてキャッシュを事前に入力
• 階層的キャッシング:マルチレベルキャッシュ戦略(GPU、CPU、ディスク)
• 動的キャッシュサイジング:リアルタイムでキャッシュサイズを調整

実世界のユースケース

チャットボットと対話エージェント

複数ターンの会話でプロンプトプレフィックスを再利用します。KV Cacheは最初のトークンまでの時間(TTFT)と全体的なレイテンシを削減します。

コード生成と補完

コードアシスタント(Copilot)は、既存のコードコンテキストに基づいた即座の補完にKV Cacheを使用します。

カスタマーサポート自動化

コンテキスト付きチャット履歴を効率的に管理し、複数の顧客インタラクション全体で低レイテンシの応答を実現します。

ドキュメントとコンテンツ生成

長文コンテンツ作成は、プロンプトキャッシングの恩恵を受け、効率的な編集と反復的なワークフローを可能にします。

ゲームとインタラクティブストーリーテリング

ゲーム内対話エンジンは、没入感のある低レイテンシのプレイヤー体験のためにストーリーコンテキストをキャッシュします。

ケーススタディ: Anthropic ClaudeのAPIは、キャッシュされたトークンに対して10倍安い料金を請求します。カスタマーサポートチャットボットで安定したプレフィックスを維持することで、運用コストを削減し、応答性を向上させます。

制限と課題

メモリ増加: K/V Cacheはコンテキストとともに線形に増加します。非常に長いコンテキストはGPUメモリを使い果たす可能性があります。

キャッシュ無効化: 以前のトークンへの変更(プロンプト編集、コンテキスト変更)は、キャッシュの一部またはすべてを無効にします。

管理の複雑さ: マルチユーザー/マルチターンシステムは、慎重なキャッシュ管理を必要とします。

訓練の制限: KV Cacheは推論専用の最適化であり、訓練中は使用されません。

緩和戦略:

• キャッシュサイズ制限と削除ポリシーを実装
• メモリ使用量を監視し、アラートを実装
• 長いコンテキストにはPaged Attentionなどの高度な技術を使用
• キャッシュ効率を考慮してプロンプトを設計

よくある質問

KV Cacheは訓練中に機能しますか? いいえ。KV Cacheは推論専用の最適化です。訓練ではすべての注意を並列に計算します。

KV Cacheはどのくらいのメモリを使用しますか? メモリはシーケンス長とモデルサイズに比例して線形にスケールします。長いシーケンスの場合、キャッシュは大量のGPUメモリを消費する可能性があります。

KV Cacheはストリーミング応答で使用できますか? はい。KV Cacheは特にストリーミングに効果的です。各新しいトークンが以前のトークンからのキャッシュされた計算を活用するためです。

プロンプトが変更されるとどうなりますか? 変更があると、その時点からキャッシュが無効になります。最大の効率を得るために、プロンプトプレフィックスを安定させてください。

すべてのLLM APIはKV Cacheを使用していますか? ほとんどの本番環境のLLM APIは自動的にKV Cacheを使用します。具体的な実装の詳細と価格については、プロバイダーのドキュメントを確認してください。

参考文献

• Hugging Face: KV Caching Explained
• Sebastian Raschka: Understanding and Coding the KV Cache in LLMs
• Neptune: Transformers Key-Value Caching Explained
• Data Science Dojo: Unlocking the Power of KV Cache
• Kapil Sharma: KV Caching Illustrated
• Akira AI: KV Caches and Time-to-First-Token
• Andrew Szot: KV Cache
• Transformers Documentation: Generation Strategies - KV Caching
• Joaolages: KV Caching Explained (Medium)