言語モデルのメモリ削減:節約量は論文が示唆するほど大きくない

言語モデルが長いプロンプトを読むとき、各トークンのキーと値を保存します。これがKVキャッシュです。32,000トークンでは、テストしたモデルで1クエリあたり2.4GBを占めます。多くの論文は大半を捨てて重要なトークンだけ残すことを提案しています。無料GPUでこれをテストし、論文がめったに一緒に問わない2つの質問をしました。モデルは必要な事実を見つけられるか?そしてメモリは本当に戻るのか?

簡潔な答え:キャッシュの25%と12.5%では、試したすべての手法が必要な事実をほぼ毎回失いました。

無料で利用できるKaggle T4上でQwen3-4B-Instructを使いました。タスクは人工的に作成したもので、無関係な文書の中に複数のプロジェクト記録（「Project X has status S and codeword C」）を埋め込み、そのうち1件、2件、または3件すべてについて質問しました。コードワードの完全一致で採点しました。プロンプトは8k、16k、32kトークンで、各6ケースを予定しました。合意した制限時間で実行を終了したため、32kの条件では6ケース中3ケースのみ完了しています。予定した306件のうち、255件の結果を記録しました。以下の各選択手法と保持容量の組み合わせは、それぞれ15問に基づきます。

保持対象の選び方を6つ比較:削除なし、スライディングウィンドウ、ランダムトークン、key-normスコア、SnapKVとH2Oのスコアリングアイデアの独自再実装。それぞれトークンレベルとページレベル(16トークンのブロック)で実行。

出典:実行の記録結果255件。12.5%の予算ではさらに悪く、すべての手法で15問中0または1問のみ正解。

キャッシュをすべて保持した場合は15問すべてに正答しました。4分の1だけ保持した場合、最良の手法でも正答は15問中2問でした。典型的な誤答ではプロジェクト名は保持されていましたが、数字は作り出されていました。例えば「cobalt-2609」ではなく「cobalt-381」と回答しました。私たちは、スコアリングがトークンの意味を保つ一方で、答えを担うまれな数字を捨ててしまうと解釈しています。これは今回の実行から得た仮説であり、検証された仕組みではありません。

推論サーバーはキャッシュを16トークン単位のページで保存します。ページを再利用できるのは、その16トークンすべてが削除された場合だけです。離れた位置のトークンを削除しても、実際に解放される容量は想定より少なくなります。同じ実行記録を使い、このメモリ管理をシミュレーションしました。ページ単位で選択する手法はほぼすべてを解放しましたが、答えも失いました。わずかな正答率を維持した手法では、解放できた容量ははるかに少なくなりました。

出典：記録した実行に対する16トークン単位のページ管理のシミュレーション（各手法30ケースの平均）。実際のページ管理型サーバーでの測定ではありません。各スコアリング手法は論理上、1問あたり約1962MBを削除したため、320MBの解放はその16%に相当します。スライディングウィンドウ手法では論理上2376MBを削除し、そのうち2373MBがページ全体の解放に相当しました（グラフの値）。他の手法より削除量が多い理由は確認していないため、比較には注意が必要です。

もう1つの結果。使用したHugging Faceの静的キャッシュでは、削除後もデコードは速くならなかった(トークンあたり約1510〜1640ms)。これはライブラリが依然として全幅に対してアテンションを計算していることを示す。ページングサーバーが同じ挙動をするとは言えない。

• 無料のT4上の4Bオープンモデル。フロンティアモデルについては何も言えません。
• ここでのSnapKVとH2Oは独自の再実装であり、公開された数値との照合は行っていません。
• 解放ページはシミュレーションです。GPUで測定したのはキャッシュ割り当てのみです。
• 32k層は半分しか完了していません。完了した32kケースもすべて失敗しました。
• 人工的に作成した1種類のタスクだけを検証しています。実際のタスクでは、こうした情報の欠落が影響しにくい可能性があります。

これが支持するもの:きちんと検証する価値のある問い。手法はページ全体を選び、数字と名前を保持する必要がある。私たちはまだそれを作っていない。