nanoVLM における KV キャッシュ実装

Hugging Face は、簡潔な PyTorch コードベースである nanoVLM において、KV（Key-Value）キャッシングをゼロから実装しました。この最適化により、自己回帰推論プロセス中の冗長な計算を排除することで、生成速度が 38% 向上しました。

自己回帰生成における計算の冗長性

自己回帰言語モデルは、テキストをトークン単位で生成します。キャッシュを使用しない標準的なトランスフォーマー実装では、モデルは次のトークンを予測するために、これまでに生成されたすべてのトークンを含むシーケンス全体を処理しなければなりません。

トランスフォーマーは内部的に並列処理されるため、各新しいトークンの予測にはすべての層を通した完全なフォワードパスが必要です。これにより、シーケンス長に対してメモリと計算量が二次的に増加します。具体的には、モデルは各ステップで過去のすべてのトークンに対してキー ($K$) とバリュー ($V$) テンソルを再計算しますが、これらのトークンと対応する射影は変わっていません。

KV キャッシュが推論を最適化する方法

KV キャッシュは、初期プロンプトが処理された後に各層で計算されたキーとバリューを保存することで、この非効率性を緩和します。シーケンス全体を再処理する代わりに、モデルは以下の増分ワークフローに従います。

  1. 初期状態のキャッシュ: 最初のパスの後、各層の計算された $K$ と $V$ がキャッシュされます。
  2. 増分計算: その後の生成ステップでは、モデルは最新トークンに対してのみ $K$ と $V$ を計算します。
  3. キャッシュ更新: 新しい $K$ と $V$ が既存のキャッシュに追加されます。
  4. アテンション計算: 現在のトークンのクエリ ($Q$) は、キャッシュされた $K$ と $V$ と組み合わせて出力を生成します。

実際には、このキャッシュは各層ごとの辞書として保持され、"key" と "value" テンソルが形状 (batch_size, num_heads, seq_len_cached, head_dim) で格納されます。

nanoVLM における技術的実装

nanoVLM の実装は、フルシーケンスの再計算から増分更新システムへの移行を実現するために、主に 3 つのコンポーネントに対する変更を伴います。

1. アテンションブロックの更新

LanguageModelGroupedAttention クラスの forward 関数が block_kv_cache を受け取るように変更されました。キャッシュが存在する場合（モデルがプレフィルフェーズにないことを示す）、モデルは現在のトークンに対して $K_{new}$ と $V_{new}$ を計算し、キャッシュされたテンソルと連結します。キャッシュが存在しない場合は、プロンプトに対して初期計算を行います。

2. 層ごとのキャッシュ追跡

LanguageModel クラスは層ごとのキャッシュ追跡を実装しました。start_pos 引数を利用して、回転位置エンコーディングが現在の生成インデックスと正しく整合するようにし、モデルが新しく生成されたトークンのシーケンスに対する絶対位置を把握できるようにします。

3. 生成ループの分岐

VisionLanguageModel の generate() メソッドは、2 つの明確なフェーズに分割されました。

  • プレフィルフェーズ: モデルは入力プロンプト全体をエンコードし、すべての層に対して初期 KV キャッシュを構築します。
  • デコードフェーズ: モデルはトークンを順次生成し、キャッシュされたキーとバリューを使用してプロンプトや過去に生成されたトークンを再処理することなく推論を行います。

アーキテクチャ変更の概要

モジュール 元の動作 新しい動作
LanguageModelGroupedAttention.forward 各ステップで $Q$, $K$, $V$ を再計算 KV キャッシュを使用し、更新
LanguageModel.forward 前状態の記憶なし 層ごとの KV キャッシュを追跡し、start_pos を処理
VisionLanguageModel.generate 単一フェーズの生成ループ プレフィル と デコード の 2 フェーズに分割

トレードオフと影響

KV キャッシュはトークンごとの推論計算量を二次的なものから $O(\text{seq len})$ に削減し、推論を高速化し、消費者向けハードウェア上で大規模モデルを実行できるようにします。しかし、この効率性にはトレードオフがあり、キャッシュを保存するためのメモリ使用量が増加し、コードの複雑さも増します。さらに、ビームサーチのような特定の推論手法は、より複雑なキャッシュ管理が必要になるため制限される可能性があります。

Sources