MiniKV: Pushing the Limits of 2-Bit KV Cache via Compression and System Co-Design for Efficient Long Context Inference

生成日: 2026-05-24 02:20

# MiniKV: Pushing the Limits of 2-Bit KV Cache via Compression and System Co-Design for Efficient Long Context Inference
- 長文推論で大きな記憶領域を占める key-value cache（KV キャッシュ）を、2 ビット量子化と重要トークン選択で小さくする方法を扱う。
- MiniKV は、重い注意を受けるトークンと直近トークンを残し、層ごとの予算配分と専用カーネルで精度と実行効率の両立を図る。
- LongBench では 86% 前後の KV キャッシュ圧縮を行いながら精度低下を小さく保ち、単一 NVIDIA A100 で長いプロンプトと高いスループットを示した。

## 論文の面白いところ

この論文の主眼は、長文対応の大規模言語モデル（Large Language Model, LLM）を実際に動かすときの記憶量にある。モデルの重みだけでなく、推論中に各層が保持する key と value の状態、すなわち KV キャッシュが長さに比例して増えるためである。長い文書を読み、さらに複数の出力トークンを生成する場合、このキャッシュは無視できない量になる。既存の研究には、KV キャッシュを低ビットで表す量子化と、重要でなさそうなトークンを捨てる適応的 KV 削減がある。MiniKV は、この二つを単純に足せばよいとは考えない。2 ビット量子化では表現できる値の幅が狭く、トークンを選んで捨てる操作とも干渉する。さらに、注意スコアを見てトークンを選ぶ方法は、FlashAttention のように巨大な注意行列を明示的に作らない実装と相性が悪い。論文は、このアルゴリズム上の問題と実装上の問題を同じ対象として扱う。結果として、単なる圧縮率の報告ではなく、長文推論で本当に速度とメモリが改善するかを測っている。研究としては地味に見えるが、長い文脈を扱う LLM を運用する際の制約に直接触れている。

## 問題設定

LLM の自己注意では、生成の各段階で過去の key と value を参照する。これを毎回計算し直すのは高価なので、多くの実装は KV キャッシュとして保存する。入力が短い場合、この保存量はそれほど目立たない。しかし長文質問応答、長いコード補完、長い要約では、プロンプト長とバッチサイズに応じて GPU メモリを強く圧迫する。KV キャッシュを 8 ビットや 4 ビットに量子化する方法は広く使われるが、2 ビットまで下げると精度が落ちやすい。一方で、注意を多く受ける heavy hitter トークンや直近の recent window だけを残す方法もあるが、長文課題では削りすぎると性能が大きく落ちる。論文は、2 ビット量子化と適応的トークン選択を組み合わせ、同じメモリ予算の下で精度と推論速度を保つ問題を扱う。形式的には、各層の key と value に対して、量子化器と選択方針を定め、元の注意出力との差を小さくすることが目標になる。ここで重要なのは、計算上うまく見える方法でも、実際の GPU カーネルで遅くなれば意味が薄いという点である。したがって、この論文の問題設定は、モデル精度だけでなく、メモリ使用量、レイテンシ、スループットを同時に含んでいる。

## 提案手法

MiniKV は、prefill 段階で入力全体を処理し、その時点で残すべき KV 状態を選ぶ。選択には累積注意スコアを用い、強く参照される heavy hitter と、生成直前に近い recent window の双方を保持する。著者らは、長文では heavy hitter だけ、または recent window だけでは失敗する課題があると観察している。そのため、標準設定ではそれぞれ 25% ずつの予算を用いる。key はチャネル方向の小さなグループごとに量子化し、value はトークン方向に量子化する。これは、key にはチャネル方向の外れ値があり、単純なトークン単位量子化では 2 ビット時の誤差が大きくなるためである。さらに、低い層に多くの KV 予算を割り、高い層では少なくする Pyramid 方式も用いる。注意スコアを得るためには注意行列を見たいが、そのままでは長さの二乗のメモリが必要になる。MiniKV は Triton による二段階の selective flash-attention カーネルを作り、通常の注意出力と累積注意スコアを線形メモリで得る。生成段階では、2 ビットに圧縮した KV を展開しながら行列積に使う処理を融合し、カーネル呼び出しとメモリアクセスを減らす。

## 結果

評価は LongBench を中心に、LLaMA2-7B-chat、LLaMA2-13B-chat、Mistral-7B-Instruct-v0.2 などで行われている。比較対象には、FullKV、KIVI、H2O、SnapKV、Q-Hitter が含まれる。LongBench では、同程度の KV キャッシュサイズで比べた場合、MiniKV は H2O、SnapKV、Q-Hitter より高い平均性能を示した。LLaMA2-7B-chat では、MiniKV Pyramid が平均 34.65 を得て、Full Model の 35.19 に近い値を保った。論文は、MiniKV が 86% 程度の KV キャッシュ圧縮を行いながら LongBench の精度を大きく崩さないと報告している。InfiniteBench の一部課題でも、Llama3-8B-instruct で Full Model や KIVI に近い平均値を示した。ただし、短いが推論能力を要する GSM8K では、精度を保つために 90% 程度の大きな適応的 KV 予算が必要だった。これは、長文検索と数理推論では、捨ててもよい状態の性質が異なることを示している。システム面では、単一 NVIDIA A100 で、長い系列ほど MiniKV のレイテンシとスループットの利点が目立つ。最大プロンプト長は強いベースラインである KIVI より約 10% 長くなり、MiniKV の専用カーネルは標準注意実装より低いメモリ使用量で長い prefill を処理した。

## 具体例

たとえば、あるシステムが 4 万トークンほどの社内仕様書を読み、その末尾で「決済エラー E-714 の再試行条件は何か」と質問される場合を考える。通常の LLM は、仕様書全体を prefill で処理し、各層の key と value を KV キャッシュに保存する。MiniKV はこの段階で、質問や後続生成で参照されやすい heavy hitter トークンと、入力末尾に近い recent window のトークンを選ぶ。選ばれた key と value は 2 ビット表現に圧縮され、層によって残す量も変えられる。生成時には、新しい query が圧縮済みの key と照合され、関連する value から回答に必要な表現を作る。期待される出力は、「E-714 は初回失敗から 30 秒後に一度だけ再試行し、同じ応答コードが返った場合は手動確認に送る」といった、文書中にある条件の要約である。間違えやすい点は、似たエラー番号 E-741 の規則を拾うことや、文書の末尾にある一般的な再試行規則だけを答えて例外条件を落とすことである。heavy hitter だけを残すと直近の補足を失う場合があり、recent window だけを残すと文書前半の定義を失う場合がある。MiniKV が両方を残すのは、このような長文入力では重要な情報が一か所にまとまっているとは限らないためである。

MiniKV: Pushing the Limits of 2-Bit KV Cache via Compression and System Co-Design for Efficient Long Context Inference

長文推論で大きな記憶領域を占める key-value cache（KV キャッシュ）を、2 ビット量子化と重要トークン選択で小さくする方法を扱う。
MiniKV は、重い注意を受けるトークンと直近トークンを残し、層ごとの予算配分と専用カーネルで精度と実行効率の両立を図る。
LongBench では 86% 前後の KV キャッシュ圧縮を行いながら精度低下を小さく保ち、単一 NVIDIA A100 で長いプロンプトと高いスループットを示した。

論文の面白いところ

この論文の主眼は、長文対応の大規模言語モデル（Large Language Model, LLM）を実際に動かすときの記憶量にある。モデルの重みだけでなく、推論中に各層が保持する key と value の状態、すなわち KV キャッシュが長さに比例して増えるためである。長い文書を読み、さらに複数の出力トークンを生成する場合、このキャッシュは無視できない量になる。既存の研究には、KV キャッシュを低ビットで表す量子化と、重要でなさそうなトークンを捨てる適応的 KV 削減がある。MiniKV は、この二つを単純に足せばよいとは考えない。2 ビット量子化では表現できる値の幅が狭く、トークンを選んで捨てる操作とも干渉する。さらに、注意スコアを見てトークンを選ぶ方法は、FlashAttention のように巨大な注意行列を明示的に作らない実装と相性が悪い。論文は、このアルゴリズム上の問題と実装上の問題を同じ対象として扱う。結果として、単なる圧縮率の報告ではなく、長文推論で本当に速度とメモリが改善するかを測っている。研究としては地味に見えるが、長い文脈を扱う LLM を運用する際の制約に直接触れている。

問題設定

LLM の自己注意では、生成の各段階で過去の key と value を参照する。これを毎回計算し直すのは高価なので、多くの実装は KV キャッシュとして保存する。入力が短い場合、この保存量はそれほど目立たない。しかし長文質問応答、長いコード補完、長い要約では、プロンプト長とバッチサイズに応じて GPU メモリを強く圧迫する。KV キャッシュを 8 ビットや 4 ビットに量子化する方法は広く使われるが、2 ビットまで下げると精度が落ちやすい。一方で、注意を多く受ける heavy hitter トークンや直近の recent window だけを残す方法もあるが、長文課題では削りすぎると性能が大きく落ちる。論文は、2 ビット量子化と適応的トークン選択を組み合わせ、同じメモリ予算の下で精度と推論速度を保つ問題を扱う。形式的には、各層の key と value に対して、量子化器と選択方針を定め、元の注意出力との差を小さくすることが目標になる。ここで重要なのは、計算上うまく見える方法でも、実際の GPU カーネルで遅くなれば意味が薄いという点である。したがって、この論文の問題設定は、モデル精度だけでなく、メモリ使用量、レイテンシ、スループットを同時に含んでいる。

提案手法

MiniKV は、prefill 段階で入力全体を処理し、その時点で残すべき KV 状態を選ぶ。選択には累積注意スコアを用い、強く参照される heavy hitter と、生成直前に近い recent window の双方を保持する。著者らは、長文では heavy hitter だけ、または recent window だけでは失敗する課題があると観察している。そのため、標準設定ではそれぞれ 25% ずつの予算を用いる。key はチャネル方向の小さなグループごとに量子化し、value はトークン方向に量子化する。これは、key にはチャネル方向の外れ値があり、単純なトークン単位量子化では 2 ビット時の誤差が大きくなるためである。さらに、低い層に多くの KV 予算を割り、高い層では少なくする Pyramid 方式も用いる。注意スコアを得るためには注意行列を見たいが、そのままでは長さの二乗のメモリが必要になる。MiniKV は Triton による二段階の selective flash-attention カーネルを作り、通常の注意出力と累積注意スコアを線形メモリで得る。生成段階では、2 ビットに圧縮した KV を展開しながら行列積に使う処理を融合し、カーネル呼び出しとメモリアクセスを減らす。

結果

評価は LongBench を中心に、LLaMA2-7B-chat、LLaMA2-13B-chat、Mistral-7B-Instruct-v0.2 などで行われている。比較対象には、FullKV、KIVI、H2O、SnapKV、Q-Hitter が含まれる。LongBench では、同程度の KV キャッシュサイズで比べた場合、MiniKV は H2O、SnapKV、Q-Hitter より高い平均性能を示した。LLaMA2-7B-chat では、MiniKV Pyramid が平均 34.65 を得て、Full Model の 35.19 に近い値を保った。論文は、MiniKV が 86% 程度の KV キャッシュ圧縮を行いながら LongBench の精度を大きく崩さないと報告している。InfiniteBench の一部課題でも、Llama3-8B-instruct で Full Model や KIVI に近い平均値を示した。ただし、短いが推論能力を要する GSM8K では、精度を保つために 90% 程度の大きな適応的 KV 予算が必要だった。これは、長文検索と数理推論では、捨ててもよい状態の性質が異なることを示している。システム面では、単一 NVIDIA A100 で、長い系列ほど MiniKV のレイテンシとスループットの利点が目立つ。最大プロンプト長は強いベースラインである KIVI より約 10% 長くなり、MiniKV の専用カーネルは標準注意実装より低いメモリ使用量で長い prefill を処理した。

具体例

たとえば、あるシステムが 4 万トークンほどの社内仕様書を読み、その末尾で「決済エラー E-714 の再試行条件は何か」と質問される場合を考える。通常の LLM は、仕様書全体を prefill で処理し、各層の key と value を KV キャッシュに保存する。MiniKV はこの段階で、質問や後続生成で参照されやすい heavy hitter トークンと、入力末尾に近い recent window のトークンを選ぶ。選ばれた key と value は 2 ビット表現に圧縮され、層によって残す量も変えられる。生成時には、新しい query が圧縮済みの key と照合され、関連する value から回答に必要な表現を作る。期待される出力は、「E-714 は初回失敗から 30 秒後に一度だけ再試行し、同じ応答コードが返った場合は手動確認に送る」といった、文書中にある条件の要約である。間違えやすい点は、似たエラー番号 E-741 の規則を拾うことや、文書の末尾にある一般的な再試行規則だけを答えて例外条件を落とすことである。heavy hitter だけを残すと直近の補足を失う場合があり、recent window だけを残すと文書前半の定義を失う場合がある。MiniKV が両方を残すのは、このような長文入力では重要な情報が一か所にまとまっているとは限らないためである。