CacheBlendとは?RAGのKVキャッシュを融合して初回応答を高速化する技術

CacheBlendは、RAGで取得した複数チャンクのKVキャッシュを再利用しつつ、文脈間の関係に必要な一部だけを再計算するLLM推論技術です。仕組み、性能、RAG開発への活かし方を解説します。

参考文献

CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion

Jiayi Yao, Hanchen Li, Yuhan Liu, Siddhant Ray, Yihua Cheng, Qizheng Zhang, Kuntai Du, Shan Lu, Junchen Jiang

論文を見る

今回の論文

今回取り上げるのは、Jiayi Yaoらによる「CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion」です。2024年にarXivで公開された論文で、URLは https://arxiv.org/abs/2405.16444 、研究分野はRAG向けLLM推論・KVキャッシュ管理です。

選んだ理由は、RAGの遅さを「検索が遅いから」だけでなく、取得した文書をLLMに読み込ませるprefillの問題として捉え直しているからです。同じ社内規程や製品資料が何度も検索結果に現れるなら、その計算結果を賢く再利用したいところです。CacheBlendは、再利用の速さと複数文書を横断して理解する品質を両立しようとします。

どんな技術か

CacheBlendは、RAGで取得した複数の文書チャンクのKVキャッシュを融合する技術です。文書ごとにあらかじめ計算したKVキャッシュをそのままつなぐのではなく、文書間の関係に影響を受けやすいトークンだけを各層で選んで再計算します。

これにより、毎回すべての文書を読み直すfull prefillに近い回答品質を保ちながら、初回トークンまでの時間(TTFT)を短縮します。論文ではfull prefillと比べ、TTFTを2.2〜3.3倍短縮、スループットを2.8〜5倍に改善したと報告しています。

課題

RAGでは、質問に応じて複数の文書断片を取得し、プロンプトに並べてLLMへ渡します。根拠を増やせる一方、入力が長くなるほど生成前のprefillが重くなります。とくに社内検索やサポートBotでは、同じ規程・FAQ・商品仕様が多数の質問で繰り返し使われます。

  • prefix cacheは入力先頭にある共通部分だけを安全に再利用できます。しかしRAGで2番目以降に置かれるチャンクは前に並ぶ文書によって表現が変わるため、通常は再利用できません。
  • 各チャンクのKVキャッシュを位置だけ補正してつなぐfull KV reuseは高速ですが、チャンク間のcross-attentionを計算していません。たとえば一つ目の資料の数値と二つ目の資料の条件を照合する質問で、根拠同士を結び付けられず品質が落ちます。
  • 毎回full prefillを行えば品質は保てますが、再利用できる文書にもGPU計算を繰り返します。長い入力と同時利用が増えるほどTTFTとGPUコストが悪化します。

つまり実運用では、「キャッシュを再利用するなら品質が下がる」「品質を守るなら毎回遅い」という二択になりがちです。CacheBlendは、全トークンを均等に再計算しなくても、重要な相互参照だけを更新すればよいのではないかという問題設定です。

用語解説

KVキャッシュ
Transformerが過去トークンから作ったKeyとValueを保存した状態です。生成時に過去を再計算しないための仕組みですが、本論文では文書チャンク単位で事前保存し、再利用可能にします。
prefill
入力プロンプト全体を一度処理してKVキャッシュを作る段階です。RAGでは取得文書が長いほど重く、ここがTTFTを左右します。CacheBlendが削る対象は主にこの計算です。
cross-attention
ここでは、あるチャンクのトークンが前に置かれた別チャンクのトークンへ向けるattentionを指します。複数資料の比較・照合には重要ですが、チャンクを単独で事前計算すると失われます。
TTFT
Time To First Token、要求送信から最初の出力トークンが返るまでの時間です。ユーザー体験では回答全体の生成速度と同じくらい重要で、RAGの長いprefillが主な悪化要因になります。
attentionの疎性
各トークンがすべての過去トークンを同じ強さで参照するのではなく、少数の相手に大きく注意を向ける性質です。CacheBlendはこの性質を利用し、更新すべきトークンを絞ります。

技術の仕組み

CacheBlendの中心は、個別チャンクで事前計算したKVを土台にし、実際の並び順で欠けた情報だけを補う「selective KV recompute」です。単純なキャッシュ連結ではなく、Transformerの層を順に通しながら一部だけ更新します。

チャンク単位でKVを保存する

まず文書をチャンクに分け、各チャンクを単独でprefillしてKVキャッシュを保存します。質問時はretrieverが選んだチャンクをハッシュで引き当て、GPU、CPUメモリ、SSDなどからロードします。位置埋め込みは実際の入力位置に合うよう補正します。

この時点でKVを単純連結すれば高速ですが、後ろのチャンクは前のチャンクを見ていない状態です。そこでCacheBlendは全層で「どのトークンのKVが本来のfull prefillからずれそうか」を推定します。

高いずれを持つトークンだけを再計算する

理想的には、full prefillとの差が大きいHigh-KV-Deviation(HKVD)トークンを選べばよいのですが、理想値を知るにはfull prefillが必要になってしまいます。CacheBlendは、前層で観測したattentionのずれを手掛かりに次層の候補を絞ります。

各層では、選ばれたトークンだけからQ・K・Vを計算します。そのK・Vを、選ばれなかったトークンのキャッシュ済みK・Vと合わせ、選ばれたトークンが全入力を参照できるattentionを実行します。結果として、重要なトークンの表現は前のチャンクを踏まえた値に更新され、その他はキャッシュを維持します。再計算率をr%にすれば、概念的な計算量もfull prefillの約r%になります。

論文の観測では、KVのずれが大きいトークンは少数に集中します。10〜20%程度のトークンを選ぶだけでattentionのずれを大きく抑えられ、既定値として15%の再計算率を採用しています。

ロードと再計算を重ねる

再計算そのものが増えると遅くなります。CacheBlendは、層iの一部トークンを更新する計算と、次の層i+1のKVキャッシュをGPUへロードするI/Oをパイプライン化します。

論文のLlama-7B・4K文脈の例では、15%のトークン再計算は層当たり3ms、NVMe SSDから1層分のKVを読む時間は16msでした。計算をロードの待ち時間に隠せるため、より容量の大きい遅いストレージにもキャッシュを置けます。これはGPUメモリだけにキャッシュを置くより、再利用可能な文書を増やせる工夫です。

実験と結果

論文はLlama-2 7B、Llama-2 13B、Yi-34Bの三つのオープンモデルで、2WikiMQA、Musique、NarrativeQA、QMSumの四データセットを評価しました。複数根拠を必要とする質問応答、長文質問応答、会議要約を含めており、単一文書だけでは済まないRAGを想定しています。

比較方法

比較対象は、全入力を毎回計算するfull KV recompute、先頭だけ再利用するprefix caching、非先頭チャンクもそのままつなぐfull KV reuseです。回答品質はタスクに応じてF1またはROUGE-L、性能はTTFTとスループットで測定しました。CacheBlendは、検索済みチャンクのキャッシュが存在する設定で評価されています。

品質を保ったままTTFTを短縮

CacheBlendはfull KV recomputeと比べてTTFTを2.2〜3.3倍短縮し、スループットを2.8〜5倍に向上させました。重要なのは、単に速いfull KV reuseより品質面で有利だった点です。full KV reuseはチャンク数が増えるほどcross-attentionを欠く影響が大きくなりますが、CacheBlendは選択的な更新でこの欠落を補います。

再計算率を5〜18%に変えた評価では、full prefillとの品質差はF1またはROUGE-Lで最大0.002でした。値はデータセットとモデルに依存しますが、「全チャンクを再計算しないと品質を保てない」という前提を崩した結果です。

ストレージ階層の効果

KVをRAMやSSDに置く設定でも、CacheBlendは品質の低下を抑えながらTTFTを改善しました。これは再計算とロードを重ねる設計によるものです。ただし効果は、モデル、文脈長、ストレージ帯域、再計算率の関係に依存します。SSDなら常に無料で速くなる、という意味ではありません。

何に使える?

CacheBlendは、同じ知識チャンクが異なる順序・組み合わせで何度も使われるRAGほど相性がよい技術です。

社内ナレッジ検索

就業規則、製品仕様、障害対応手順などは多くの質問で再登場します。チャンク単位のKVを保存しておけば、毎回同じ長文をprefillする負担を減らせます。複数部署の規程を照合する質問でも、必要箇所を選択的に更新する考え方が役立ちます。

カスタマーサポートCopilot

FAQ、契約条件、注文状態、製品マニュアルを混ぜて回答する場合、回答の初動が遅いとオペレーター支援の価値が下がります。頻出資料のキャッシュ再利用はTTFTを下げる候補になります。ただし、リアルタイムの注文情報はキャッシュ対象にせず、鮮度が必要なデータを区別する設計が必要です。

エージェントのツール結果統合

検索、データベース照会、APIレスポンスを複数集めるエージェントでも、以前に見た定型資料を再利用できます。CacheBlend自体はエージェント制御技術ではありませんが、「各ツール結果をどの単位で再利用可能な文脈にするか」を考える基盤になります。

開発や事業へのヒント

この論文の実務的な示唆は、RAGキャッシュを検索結果のキャッシュだけで終わらせないことです。どの文書を取得したかだけでなく、LLMが読むために作った中間状態も、再利用の対象になり得ます。

文書チャンクを安定した単位で設計する

KV再利用は、同じ内容が同じトークン列として再び現れるほど効きます。チャンク分割、見出しの付け方、テンプレートを頻繁に変えるとキャッシュヒット率が落ちます。頻出の規程・仕様は安定したIDと内容ハッシュで管理し、更新時に確実に無効化する設計が重要です。

品質KPIと速度KPIを分けて測る

「キャッシュがヒットした」だけでは不十分です。複数文書を比較する代表質問でF1や人手評価を測り、TTFT、GPU使用率、再計算率と一緒に追ってください。full KV reuseのように速くても、チャンク間推論が必要な質問だけ壊れるケースを見逃しにくくなります。

小規模プロダクトでは段階的に採用する

独自のKV融合カーネルをすぐ実装する必要はありません。まずはprefix cache、頻出プロンプトの固定、取得チャンク数の制御、回答前のrerankで入力長を減らす方法が現実的です。高いトラフィックで長いRAG文脈が支配的になった段階で、CacheBlendのような非prefixキャッシュ再利用を検討する順序がよいでしょう。

限界

CacheBlendは、キャッシュ済みのチャンクが再利用されることを前提とします。質問ごとに新しい文書ばかり取得される用途では、事前計算・保存のコストを回収しにくくなります。またチャンク内容が変わればKVキャッシュは無効化しなければならず、文書更新の鮮度管理が必要です。

選択するトークンは完全なfull prefillとの差を直接見ず、前層のattention情報から近似します。論文の評価では小さな品質差に収まりましたが、モデル、プロンプト、文書ドメイン、必要な推論の種類が変われば再計算率の再調整が必要です。特に表・コード・多言語資料などで同じ品質が出るとは限りません。

さらに、実装は層ごとのKVロード、部分attention、GPU/CPU/SSDの転送を制御する推論基盤に踏み込みます。クラウドLLM APIだけを利用するアプリでは直接導入できません。保存されたKVには元文書由来の情報が含まれるため、テナント分離、暗号化、アクセス制御、無効化も通常の検索キャッシュ以上に慎重に設計する必要があります。

よくある質問

Q. CacheBlendは通常のprefix cacheを置き換えますか?

A. 置き換えではなく補完です。先頭に共通するシステムプロンプトや固定資料はprefix cacheで安全に再利用できます。CacheBlendの価値は、RAGの2番目以降のように非prefix位置に置かれ、前の文書との関係で表現が変わるチャンクにも再利用の余地を作る点です。

Q. なぜKVキャッシュを単純に連結すると回答品質が下がるのですか?

A. チャンクを単独で計算したKVには、実際には前に並ぶ別チャンクへのattentionが入っていません。二つの資料の数字や条件を結び付ける質問では、そのcross-attentionが必要です。CacheBlendは影響の大きいトークンだけを再計算して、この不足を補います。

Q. 15%だけの再計算で常に十分ですか?

A. いいえ。論文では15%を既定値とし、5〜18%の範囲で小さな品質差を示しましたが、最適値はモデル、文書数、文脈長、タスクに依存します。導入時は品質評価セットを作り、TTFTとのトレードオフを測って決める必要があります。

Q. RAGの検索結果キャッシュと何が違いますか?

A. 検索結果キャッシュは「どの文書を返すか」を再利用します。CacheBlendは、その文書をLLMが読むために計算したTransformer内部のKV状態を再利用します。両者を組み合わせると速くなり得ますが、検索結果が古い場合はKVも使ってはいけません。

Q. SaaSで導入する際の最初の確認事項は何ですか?

A. 頻出する文書チャンクの割合、要求ごとの入力長、TTFT、文書更新頻度を測ることです。ヒット率が低い、文書が頻繁に更新される、入力が短い場合は、KV融合の複雑さより単純な入力削減やprefix cacheの方が費用対効果に優れる可能性があります。

今日の学び

この論文は、RAGで複数の取得文書を読むたびに長いprefillを繰り返す課題を扱いました。CacheBlendは、チャンクごとのKVキャッシュを再利用しながら、文書間のattentionに重要な一部トークンだけを再計算して解こうとします。開発では、検索結果だけでなくLLMの中間計算も再利用対象として捉え、品質・TTFT・キャッシュ鮮度をセットで最適化することがヒントになります。