2025-11-16T09:58:12.370377

Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference

Feng, Lv, Cao et al.

Large Language Models have excelled in various domains but face efficiency challenges due to the growing Key-Value (KV) cache required for long-sequence inference. Recent efforts aim to reduce KV cache size by evicting vast non-critical cache elements during runtime while preserving generation quality. However, these methods typically allocate compression budgets uniformly across all attention heads, ignoring the unique attention patterns of each head. In this paper, we establish a theoretical loss upper bound between pre- and post-eviction attention output, explaining the optimization target of prior cache eviction methods, while guiding the optimization of adaptive budget allocation. Base on this, we propose {\it Ada-KV}, the first head-wise adaptive budget allocation strategy. It offers plug-and-play benefits, enabling seamless integration with prior cache eviction methods. Extensive evaluations on 13 datasets from Ruler and 16 datasets from LongBench, all conducted under both question-aware and question-agnostic scenarios, demonstrate substantial quality improvements over existing methods. Our code is available at https://github.com/FFY0/AdaKV.

academic

Ada-KV: 効率的なLLM推論のための適応的予算配分によるKVキャッシュ削除の最適化

基本情報

論文ID: 2407.11550
タイトル: Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference
著者: Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, S. Kevin Zhou
分類: cs.CL cs.AI
発表時期/会議: 第39回ニューラル情報処理システム会議(NeurIPS 2025)
論文リンク: https://arxiv.org/abs/2407.11550

要約

大規模言語モデル(LLMs)は様々な領域で優れた性能を発揮していますが、長系列推論における増加し続けるKey-Value(KV)キャッシュの需要により、効率の課題に直面しています。最近の研究では、実行時に大量の非重要キャッシュ要素を削除することでKVキャッシュサイズを削減しながら、生成品質を維持しています。しかし、これらの方法は通常、すべての注意ヘッド間で圧縮予算を均等に配分し、各ヘッドの独特な注意パターンを無視しています。本論文は、削除前後の注意出力間の理論的損失上界を確立し、先行するキャッシュ削除方法の最適化目標を説明しながら、適応的予算配分の最適化を指導しています。これに基づいて、著者らはAda-KVを提案しました。これは初めてのヘッドレベル適応的予算配分戦略です。本方法はプラグアンドプレイの利点を有し、既存のキャッシュ削除方法とシームレスに統合できます。

研究背景と動機

問題記述

大規模言語モデルが処理する系列長の継続的な増加(例えば、GPTは128K、Claude3は200K、Gemini-Pro-1.5は2M トークンをサポート)に伴い、KVキャッシュのメモリ需要は指数関数的に増加しています。8Bパラメータを持つLLMの場合、単一の2M トークン系列を処理するには最大256GBのキャッシュが必要になる可能性があり、GPUメモリ効率と計算実行時効率に深刻な影響を与えます。

既存方法の限界

既存のキャッシュ削除方法は主に2つのカテゴリに分類されます:

スライディングウィンドウ削除方法: 初期および最新のキャッシュ要素を単純に保持しますが、生成品質を大幅に低下させます
Top-k削除方法: 注意重みに基づいて重要なキャッシュ要素を選択しますが、すべての注意ヘッド間で予算を均等に配分します

重要な問題は、既存の方法が異なる注意ヘッドの独特な特性を無視していることです。一部のヘッドは疎な注意集中パターンを持ち、他のヘッドの注意分布はより分散しています。

研究動機

Llama-3.1-8B-Instructモデルを分析することで、著者らは大多数の注意ヘッドが非常に小さなキャッシュ比率(例えば、上位5%)のみで、ほぼすべての注意重みを保持できることを発見しました。一方、分散したヘッドはより大きなキャッシュ比率を必要とします。この不均一な注意集中パターンは、適応的予算配分の理論的基礎を提供します。

核心的貢献

適応的予算配分戦略: 各注意ヘッドの独特な注意パターンに応じて予算配分を動的に調整できる、初めてのヘッドレベル適応的予算配分戦略Ada-KVを提案しました
理論的枠組みの確立: キャッシュ削除の理論的枠組みを確立し、削除損失を定義し、その上界を導出し、既存方法の最適化目標を説明し、Ada-KVの設計を指導しています
プラグアンドプレイ互換性: Ada-KVはプラグアンドプレイ特性を持ち、既存のキャッシュ削除方法にシームレスに統合でき、効率的なCUDAカーネル実装により計算効率を維持しています
包括的な実験検証: RulerおよびLongBenchの29個のデータセットで包括的な評価を実施し、質問認識および質問非認識の両シナリオで顕著な改善を示しています

方法の詳細説明

タスク定義

マルチヘッド自己注意層が与えられた場合、予算制約下で保持するKVキャッシュ要素を選択し、削除後の注意出力と元の出力間の損失を最小化します。

理論的基礎

L1削除損失の定義

著者らは削除損失を、削除前後の自己注意機構出力間のL1距離として定量化しました:

$\text{L1 Eviction Loss} = ||y - \hat{y}||_1$

ここで、 $y$ と $\hat{y}$ はそれぞれ削除前後の注意出力です。

損失上界の導出

定理3.1: L1削除損失は $\epsilon$ 上界で制限できます:

$\text{L1 Eviction Loss} \leq \epsilon = 2hC - 2C\sum_{i \in [1,h]}\sum_{j \in [1,n]} I_i^j A_i^j$

ここで、 $C = \max\{\|V_iW_i^O\|_\infty\}$ は定数、 $I_i^j$ は削除決定指示変数、 $A_i^j$ は注意重みです。

定理3.2: Top-kキャッシュ削除方法は与えられた予算配分下で損失上界を最小化できます:

$\epsilon^* = 2hC - 2C\sum_{i \in [1,h]}\sum_{A_i^j \in \text{Top-k}(A_i, k=B_i)} A_i^j$

Ada-KVアルゴリズム

アルゴリズム1: 適応的予算配分

入力: 総予算B、各ヘッド注意重み{A_i}
出力: 配分予算{B_i^*}
1. すべてのヘッドの注意重みを連結: A = Cat({A_i})
2. Aから上位B個の重みを選択: Top-k(A, k=B)
3. 各ヘッドで選択された重みの数を統計: {f_i}
4. 配分予算を設定: {B_i^* = f_i}