NVIDIA KVPress: 長文コンテキストLLM向けKVキャッシュ圧縮ツールキット

NVIDIAはKVPressを導入しました。これは、Large Language Models (LLM) のKey-Value (KV) キャッシュを圧縮するように設計されたPythonツールキットで、長いコンテキストウィンドウの効率的な処理を可能にします。中間注意結果の保存に関連するメモリオーバーヘッドを削減することにより、KVPressはGPUメモリを過剰に必要とすることなく、インコンテキスト検索や拡張推論に必要なようなより長いシーケンスをモデルが扱えるようにします。

KVキャッシュメモリボトルネック

自己回帰LLMでは、テキスト生成はトークンごとに行われます。以前のすべてのトークンの表現を新しいトークンを生成するたびに再計算するのを避けるため、モデルは注意層からのキー(K)と値(V)を保存するKVキャッシュを使用します。これにより計算が最適化されますが、キャッシュサイズはコンテキストウィンドウに対して線形に増加します。

KVキャッシュのメモリ消費は次の式で決定されます:

$$ ext{Size} \left( ext{KV} ight) = 2 imes ext{precision} imes n_{layers} imes n_{heads} imes d imes n_{tokens}$$

bfloat16精度で動作するLlama 3-70Bのようなモデルでは、1Mトークンのコンテキストにおいて、KVキャッシュだけで約327.6 GBが必要です。モデルウェイトに必要な140 GBを加えると、総メモリ需要は約470 GBに達し、これはKVキャッシュが総メモリ使用量の約70%を占めていることを意味します。

KVPress: モジュラー圧縮ツールキット

KVPressは、「プレス」と呼ばれる圧縮アルゴリズムのスイートを提供します。これらは、メモリフットプリントを削減するために重要度の低いKVペアを剪定します。これらのプレスはフォワードフックを介してモデルの注意層に統合され、プリフィリングフェーズを対象とし、キャッシュが最大になるときに圧縮します。

利用可能な圧縮技術

KVPressはいくつかの最先端の剪定手法を実装しており、以下を含みます:

• KnormPress: キー・バリューのノルムが最も低いKVペアを剪定します。
• SnapKVPress: 最新のクエリに対する注意重みが低いKVペアを剪定します。
• ExpectedAttentionPress: KVPressの作者によって作成された新しい未公開の手法で、将来のクエリに対する期待される注意重みが最も低いKVペアを剪定します。

KVPressはモジュラー設計になっており、研究者は新しい手法を簡単に追加してツールキットを拡張し、transformersライブラリで利用可能なKVキャッシュ量子化などの他のメモリ節約技術と統合することができます。

パフォーマンスとメモリへの影響

KVキャッシュを圧縮すると、ピークメモリ使用量が削減され、生成速度が向上します。bfloat16精度で128kトークンのコンテキスト長を持つLlama 3.1 8Bに、50%の圧縮率でKVPressを適用すると、以下の改善が得られます:

• メモリ削減: ピークメモリ使用量が45GBから37GBに低下。
• 速度向上: A100 GPU上でデコード速度が11トークン/秒から17トークン/秒に増加。

ベンチマークと精度のトレードオフ

KVPressには、RULER、InfiniteBench、Loogleなどの標準的な長文コンテキストデータセット上で圧縮手法をベンチマークするCLIが含まれています。

RULERデータセット（4kトークンのコンテキスト長）でのベンチマークでは、最も効果的な組み合わせはAdaKVPressとExpectedAttentionPressであることが示されています。ただし、ツールキットは重要なトレードオフを指摘しています：圧縮率が上がるにつれて、モデルの精度は一般的に低下します。これにより、モデルのパフォーマンスへの影響を最小限に抑える圧縮アルゴリズムの研究が継続的に必要であることが強調されます。