同じモデルなのに、会話が長くなると VRAM が増える理由

KV cache は、注意機構の計算に必要なキーとバリューを保存し、後続の token を生成するときの重複計算を減らします。Hugging Face のキャッシュ戦略のドキュメントでは、各方式のメモリと速度のトレードオフを説明しています。NVIDIA の推論に関する記事でも、キャッシュをリソース予算の重要な要素に挙げています。Hugging Face ドキュメント ↗、NVIDIA 技術記事 ↗

一般的なフルアテンションのキャッシュを例にすると、簡略化した見積もりは次のとおりです。

キャッシュのバイト数 ≈ 2 × 層数 × KV ヘッド数 × 各ヘッドの次元数 × token 数 × 同時実行するシーケンス数 × 要素あたりのバイト数。

ここでの 2 は、キーとバリューに対応します。例として、32 層、KV ヘッド 8 個、各ヘッド 128 次元のモデルを考えます。4,096 token を保存し、単一シーケンスで各値に 2 バイトを使うと、キャッシュは約 512MiB になります。

これらの仮定が変わらなければ、次のようになります。

• コンテキストが 8,192 token になると、約 1GiB。
• 4,096-token のシーケンスを同時に 4 本保存すると、約 2GiB。

これは規模感をつかむために独自に作った計算例で、キャッシュの割り当てやバッファなどのオーバーヘッドを含んでいません。スライディングウィンドウ、異なる注意機構の構造、キャッシュの量子化によっても結果が変わるため、すべてのモデルに当てはめることはできません。

試してみる

モデルと量子化方式を固定し、入力長を段階的に増やして、メモリのピーク使用量と最初のトークンが出るまでの時間を記録してください。その後、同時実行数を変えます。条件を 1 回に 1 つだけ変更すれば、どの要求がメモリ使用量を押し上げたか判断しやすくなります。

VRAM を節約する方法では、速度も測りましょう。キャッシュを別のメモリ領域に移したり、さらに量子化したりすると、新たな計算や転送のコストが生じる可能性があります。