0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

エージェントで膨らむKVキャッシュを、prefixキャッシュを壊さず刈るIntentKV

0
Posted at

ツール呼び出しを繰り返すエージェントを長く走らせていると、最初に音を上げるのはモデルの計算そのものではなく、たいていメモリのほうだ。検索ツールが返した10万件規模のドキュメント、丸ごと貼り付いたWebページのHTML、過去の推論ログ。それらがコンテキストに追記され続け、二度と外れない。この積み上がった履歴を保持するのが KVキャッシュ で、エージェントの実運用ではここがボトルネックになる。

KVキャッシュは、過去のトークンごとに計算したKey/Value(注意機構の中間結果)を保存しておく領域だ。これがあるおかげで、新しいトークンを生成するたびに全履歴を再計算せずに済む。ただしサイズはトークン数に比例して膨らむ。単発のチャットなら問題にならないが、数十ターンにわたってツール出力を飲み込むエージェントでは、パラメータ計算よりKVの読み書き帯域が先に詰まる。

6月にarXivへ出た「IntentKV」は、この問題に正面から取り組みつつ、既存のKV削減手法が見落としていた副作用を突いている。

「古いトークンを捨てる」の隠れたコスト

KVが膨らむなら古いトークンを間引けばいい、という発想の手法はすでにいくつもある。末尾の注意分布から重要トークンを拾うSnapKV、影響の大きい「ヘビーヒッター」を残すH2O、冒頭のシンク+直近ウィンドウだけ残すStreamingLLMあたりが代表格だ。

問題は、削ったあとにKVの行を詰め直す(compaction)と、トークンの位置がずれてしまう点にある。位置がずれると、RoPE(回転位置埋め込み)の位相も、キャッシュ上のスロット識別子も変わる。ここで壊れるのが prefixキャッシュ だ。vLLMのautomatic prefix cachingに代表されるこの仕組みは、複数リクエストで先頭の文脈が一致するとき、その部分のKV計算を丸ごと使い回す。エージェントは毎ターンほぼ同じ長い前提(システムプロンプト+これまでの履歴)を送り直すので、prefixキャッシュのヒット率がそのまま応答速度と課金に効く。

from vllm import LLM
# 先頭が一致するリクエスト間でKV計算を再利用する
llm = LLM(model="Qwen/Qwen2.5-14B", enable_prefix_caching=True)

つまり従来のKV削減は、メモリを節約する裏で、詰め直しによってprefixキャッシュの一致を壊し、再計算というかたちで節約分を払い戻していた。論文の計測では、compaction系のprefixヒット率は0〜3%まで落ち込む。個人的にはここが一番の勘所だと思う。KV削減とキャッシュ再利用は別々の最適化として語られがちだが、実運用では明確にトレードオフの関係にある。

詰め直さずに「無効化」する

IntentKVの一つ目のアイデアは拍子抜けするほど素朴だ。捨てたいトークンのKV行を消して詰めるのではなく、その物理スロットを1個の「番兵(sentinel)スロット」に張り替える。番兵のKeyには極端に小さい値(K = -10⁴)を入れておくので、softmaxを通すと注意の重みがゼロに潰れ、事実上そのトークンは存在しないのと同じ扱いになる。

生き残ったトークンは元の論理位置・物理スロットに居座ったままなので、RoPEの位相もスロット識別子も変わらない。結果としてprefixキャッシュの木構造(radix tree)は次ターンでもマッチし続ける。この張り替えだけで、8kトークン予算でのprefixヒット率が20.7%まで戻る。0〜3%との差は、そのままエージェントの1ターンあたり再計算量の差になる。

二つ目は、どのトークンを残すかの判断をターンをまたいで賢くする部分だ。IntentKVはセッション単位で「意図メモリ」を持ち、ターンごとのクエリを減衰させながら足し込んでいく。

M_t = e^(-λ)·M_(t-1) + Enc(q_t)   (λ=0.5)

直近のプロンプト1個だけでトークンの重要度を測るのではなく、エージェントがこのセッションで何を追い続けているかという「意図の履歴」に照らして残すトークンを選ぶ。さらに凍結したモデルの上に、ゼロ初期化した残差ヘッドを乗せて学習補正をかける。ゼロ初期化なので、学習が効くまではルールベースの素点が下限として保証される作りだ。ベースモデル自体には一切手を入れない。

同じ予算での正答率

効きどころは、同じKV予算に絞ったときの正答率だ。エージェント向けの深掘り検索ベンチ(約10万文書を対象とする830問のBrowseComp-Plus、多段QAのFRAMES)で、KV予算を8kトークンに固定して比較した結果がこれ(Qwen2.5-14B)。

手法 8k予算での正答率
IntentKV 18.55%
StreamingLLM 8.19%
SnapKV 7.71%
H2O 5.30%

絶対値が低いのはベンチ自体が難しいためで、見るべきは相対差だ。同じメモリ予算で他手法の倍以上を出している。Qwen3-8Bでは、8k予算のIntentKVが14.10%で、16k予算のフルキャッシュ(15.06%)にほぼ肉薄した。半分の予算で精度をほとんど落とさない、という主張には説得力がある。メモリ側の数字も派手で、長いクエリ100件の最悪ケースでピークトークンが92.3k→20.5k、KV読み出しが411M→31M(約92.6%減)まで下がっている。

現時点での距離感

とはいえ、明日から使えるわけではない。論文自身が限界を認めている。KV予算Cはクエリごとに動的調整されず全体固定、残差ヘッドの学習はToolBenchのマルチターン軌跡に限られ、検証は8B〜14Bまで。コードの公開も現時点では確認できない。だから今すぐ導入する話ではなく、方向性を示した研究として読むのが正しい。

ただ、番兵スロットで無効化するという発想は単純で、KV削減とprefixキャッシュを一つの問題として設計する視点は、推論スタック側が独立に取り込める余地がある。実務でエージェントをvLLMなどで回している人にとっての当面の教訓は、enable_prefix_cachingを確実に有効にしておくこと、そして積極的なKV圧縮を入れるときはprefixヒット率が静かに落ちていないかを併せて計測することだ。片方だけ見ていると、節約したつもりが再計算で取り返される。IntentKVはその落とし穴に名前を付けた論文でもある。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?