📄 論文解説: RAGCache — 知識ツリー構造によるRAG推論のKVキャッシュ最適化

本記事は RAGCache: Efficient Knowledge Caching for Retrieval-Augmented Generation の解説記事です。

論文概要（Abstract）

著者らは、RAG（Retrieval-Augmented Generation）システムにおいて、検索されたドキュメントのKVキャッシュを「知識ツリー（Knowledge Tree）」構造で階層的に管理するRAGCacheシステムを提案している。従来のRAGパイプラインでは、検索で取得したドキュメントのKVテンソルを毎リクエスト再計算していたが、RAGCacheは頻繁にアクセスされるドキュメントのKVキャッシュを保持・再利用する。著者らは、vLLMバックエンドとの統合評価で、スループット最大4倍向上・TTFT最大3.5倍削減を報告している。

この記事は Zenn記事: LangGraph×Claude Sonnet 4.6のプロンプトキャッシュ最適化でAgentic RAGコスト90%削減 の深掘りです。

情報源

• arXiv ID: 2401.02038
• URL: https://arxiv.org/abs/2401.02038
• 著者: Chaoyi Jiang, Hanqi Li, Yimin Jiang, et al.
• 発表年: 2024
• 分野: cs.CL, cs.AI

背景と動機（Background & Motivation）

RAGシステムでは、ユーザークエリに対して外部知識ベースからドキュメントを検索し、それをLLMのコンテキストに追加して回答を生成する。このプロセスでは、検索で取得したドキュメント（通常数千トークン）のKVテンソルを毎回プレフィルフェーズで計算する必要がある。

しかし、実際のRAGシステムでは同一のドキュメントが異なるクエリに対して繰り返し検索されるパターンが一般的である。例えば社内FAQシステムでは、人気の高い規約文書やガイドラインが多くのクエリで検索される。著者らは、Natural Questions・TriviaQA・HotpotQAの分析で、上位10%のドキュメントが全検索ヒットの40〜60%を占めることを確認している。

この偏在的なアクセスパターンを活用し、頻出ドキュメントのKVキャッシュを保持・再利用することで、RAG推論の効率化を図るのがRAGCacheの基本アイデアである。

主要な貢献（Key Contributions）

• 知識ツリー（Knowledge Tree）構造: ドキュメントIDをキーとするツリー構造でKVキャッシュを階層管理。共通のシステムプロンプトを根ノード、個別ドキュメントを葉ノードとして、キャッシュの共有と独立管理を両立
• 適応的キャッシュ置換ポリシー: LRU（Least Recently Used）を基本とし、ドキュメントの人気度・KVキャッシュサイズ・検索頻度を考慮した動的な置換戦略
• vLLMとの統合設計: vLLMのPagedAttentionの上位レイヤーとして機能し、既存のLLMサービングスタックに追加可能
• スループット4倍向上: Natural Questions/TriviaQA/HotpotQAでの評価で、標準RAGパイプラインと比較してスループットを最大4倍改善

技術的詳細（Technical Details）

知識ツリーの構造

著者らが提案する知識ツリーは、RAGプロンプトの構造をツリーとして表現する。

           [Root: System Prompt KV]
           /         |          \
    [Doc A KV]   [Doc B KV]   [Doc C KV]    ← 検索ドキュメントのKVキャッシュ
       |            |            |
  [Query KV]   [Query KV]   [Query KV]     ← クエリ固有部分（キャッシュ対象外）

根ノードにはシステムプロンプトのKVキャッシュが保持され、すべてのリクエストで共有される。各ドキュメントのKVキャッシュは中間ノードとして管理され、クエリが同じドキュメントセットを使用する場合にキャッシュヒットする。

キャッシュ管理の数学的定式化

著者らはキャッシュ管理を以下のように定式化している。

GPUメモリ容量 $M$ のうち、モデルパラメータに $M_{\text{model}}$、KVキャッシュに $M_{\text{kv}}$ を割り当てる。

\[M_{\text{kv}} = M - M_{\text{model}}\]

各ドキュメント $d_i$ のKVキャッシュサイズは:

\[\text{size}(d_i) = 2 \times L \times H \times d_k \times |d_i| \times \text{precision}\]

ここで、

• $L$: モデルのレイヤー数
• $H$: アテンションヘッド数
• $d_k$: ヘッドの次元数

• $

  d_i

  $: ドキュメント $d_i$ のトークン数

• $\text{precision}$: 数値精度のバイト数（FP16なら2）

キャッシュに保持可能なドキュメント数は:

\[N_{\text{cached}} = \left\lfloor \frac{M_{\text{kv}}}{\bar{s}} \right\rfloor\]

ここで $\bar{s}$ は平均ドキュメントKVキャッシュサイズである。

適応的キャッシュ置換アルゴリズム

著者らは単純なLRUではなく、以下のスコアに基づく置換ポリシーを提案している。

\[\text{score}(d_i) = \alpha \cdot \text{freq}(d_i) + \beta \cdot \text{recency}(d_i) - \gamma \cdot \frac{\text{size}(d_i)}{\bar{s}}\]

ここで、

• $\text{freq}(d_i)$: ドキュメント $d_i$ の検索頻度（正規化済み）
• $\text{recency}(d_i)$: 最後にアクセスされてからの経過時間の逆数
• $\text{size}(d_i)$: KVキャッシュサイズ（大きいドキュメントはペナルティ）
• $\alpha, \beta, \gamma$: ハイパーパラメータ（著者らのデフォルト: $\alpha=0.5, \beta=0.3, \gamma=0.2$）

スコアが最も低いドキュメントのKVキャッシュが優先的にevictされる。

アルゴリズムの擬似コード

from collections import OrderedDict
from dataclasses import dataclass, field
import torch


@dataclass
class CachedDocument:
    """キャッシュされたドキュメントのKVテンソル"""
    doc_id: str
    kv_cache: tuple[torch.Tensor, torch.Tensor]  # (K, V)
    token_count: int
    access_count: int = 0
    last_access_time: float = 0.0


class RAGCacheManager:
    """RAGCache知識ツリー管理（著者らのアルゴリズムを再現）"""

    def __init__(self, max_memory_bytes: int, alpha: float = 0.5,
                 beta: float = 0.3, gamma: float = 0.2):
        self.max_memory = max_memory_bytes
        self.used_memory = 0
        self.cache: dict[str, CachedDocument] = {}
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma

    def get_or_compute(
        self,
        doc_id: str,
        doc_tokens: list[int],
        model,
        current_time: float,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """ドキュメントのKVキャッシュを取得（キャッシュミス時は計算）

        Args:
            doc_id: ドキュメント識別子
            doc_tokens: ドキュメントのトークン列
            model: LLMモデル
            current_time: 現在時刻（recency計算用）

        Returns:
            (K, V) テンソルのタプル
        """
        if doc_id in self.cache:
            # キャッシュヒット
            entry = self.cache[doc_id]
            entry.access_count += 1
            entry.last_access_time = current_time
            return entry.kv_cache

        # キャッシュミス: KVテンソルを計算
        kv = model.compute_kv(doc_tokens)
        kv_size = self._estimate_size(len(doc_tokens))

        # メモリが不足する場合、スコアが低いエントリをevict
        while self.used_memory + kv_size > self.max_memory and self.cache:
            self._evict_lowest_score(current_time)

        # 新規エントリを追加
        self.cache[doc_id] = CachedDocument(
            doc_id=doc_id,
            kv_cache=kv,
            token_count=len(doc_tokens),
            access_count=1,
            last_access_time=current_time,
        )
        self.used_memory += kv_size

        return kv

    def _compute_score(self, entry: CachedDocument, current_time: float) -> float:
        """eviction用スコア計算（論文Section 3.2のスコア関数）"""
        freq_norm = entry.access_count / max(1, sum(
            e.access_count for e in self.cache.values()
        ))
        recency = 1.0 / max(1.0, current_time - entry.last_access_time)
        size_norm = entry.token_count / max(1, sum(
            e.token_count for e in self.cache.values()
        ) / len(self.cache))

        return self.alpha * freq_norm + self.beta * recency - self.gamma * size_norm

    def _evict_lowest_score(self, current_time: float) -> None:
        """スコアが最も低いエントリをevict"""
        if not self.cache:
            return
        worst_id = min(
            self.cache.keys(),
            key=lambda did: self._compute_score(self.cache[did], current_time),
        )
        entry = self.cache.pop(worst_id)
        self.used_memory -= self._estimate_size(entry.token_count)

    def _estimate_size(self, token_count: int) -> int:
        """KVキャッシュサイズの推定（FP16, LLaMA-7B想定）"""
        layers = 32
        heads = 32
        head_dim = 128
        precision_bytes = 2  # FP16
        return 2 * layers * heads * head_dim * token_count * precision_bytes

実験結果（Results）

著者らはNatural Questions・TriviaQA・HotpotQAの3つのQAベンチマークで評価を行っている（論文Section 5より）。

データセット	ベースライン スループット	RAGCache スループット	向上倍率
Natural Questions	12.3 req/s	48.7 req/s	3.96x
TriviaQA	10.8 req/s	38.2 req/s	3.54x
HotpotQA	8.5 req/s	28.1 req/s	3.31x

TTFTの改善（論文Table 3より）:

データセット	ベースライン TTFT	RAGCache TTFT	削減率
Natural Questions	450ms	130ms	71.1%
TriviaQA	520ms	160ms	69.2%
HotpotQA	680ms	240ms	64.7%

著者らの報告によると、キャッシュヒット率はドキュメントの人気度分布に強く依存する。Zipf分布に従うアクセスパターン（上位10%のドキュメントが全アクセスの50%以上を占める）では、GPUメモリの20%をキャッシュに割り当てるだけで80%以上のヒット率を達成している。

精度への影響

著者らは、RAGCacheが生成品質に影響を与えないことを確認している。KVキャッシュの再利用は数値的に等価であり、Exact Match（EM）スコアはベースラインと同一であると報告されている。

実装のポイント（Implementation）

vLLMとの統合

RAGCacheはvLLMのPrefix Cachingの上位レイヤーとして動作する。vLLMの--enable-prefix-cachingフラグと組み合わせることで、ドキュメントレベルとプレフィックスレベルの2層キャッシュが実現される。

キャッシュウォームアップ戦略

著者らはコールドスタート問題を認識しており、以下の対策を推奨している:

• 事前ウォームアップ: 人気ドキュメントのKVキャッシュを起動時に事前計算
• キャッシュ永続化: シャットダウン時にキャッシュをディスクに保存し、再起動時に復元

Zenn記事との関連

Zenn記事で解説されている4層キャッシュブレークポイント（ツール定義→システムプロンプト→RAGコンテキスト→会話履歴）のうち、RAGCacheは第3層（RAGコンテキスト）の最適化に直接対応する。APIレベルのプロンプトキャッシュ（Anthropicのcache_control）とRAGCacheは異なるレイヤーで動作するが、目的は同一である。

実運用への応用（Practical Applications）

社内検索システムでの活用

Zenn記事が想定する社内文書検索アシスタントは、RAGCacheの最適な適用対象である。社内FAQや規約文書は更新頻度が低く、同一ドキュメントが繰り返し検索されるため、高いキャッシュヒット率が期待できる。

キャッシュ効率が低下するケース

著者らは以下のケースでキャッシュ効率が低下することを報告している:

• ドキュメントの更新頻度が高い場合（ニュースサイト等）: KVキャッシュの無効化が頻発
• ドキュメントの人気度が均一に分散する場合: キャッシュヒット率が低下
• ドキュメント数がGPUメモリに対して大きすぎる場合: evictionが頻発

関連研究（Related Work）

• Prompt Cache (Gim et al., 2023): プロンプトモジュール単位でのKVキャッシュ再利用。RAGCacheはこの概念をRAG固有のドキュメントレベルに拡張
• vLLM / PagedAttention (Kwon et al., SOSP 2023): KVキャッシュのメモリ管理基盤。RAGCacheはvLLMの上位レイヤーとして動作
• CacheGen (Liu et al., 2023): KVキャッシュの圧縮・転送最適化。分散環境でRAGCacheと組み合わせ可能

まとめと今後の展望

RAGCacheは、RAGシステムの検索ドキュメントのKVキャッシュを知識ツリー構造で管理することで、スループット最大4倍・TTFT最大3.5倍の改善を実現している。人気度に基づく適応的キャッシュ置換により、限られたGPUメモリで高いヒット率を維持する設計が特徴である。

今後の課題として、マルチホップ検索での複合ドキュメントキャッシュ、分散環境でのキャッシュ共有が挙げられている。

参考文献

• arXiv: https://arxiv.org/abs/2401.02038
• Related Zenn article: https://zenn.dev/0h_n0/articles/555a4e799660de