Home 論文解説: CachedAttention — マルチターン会話のKVキャッシュ再利用でTTFT 2.8倍高速化
投稿
キャンセル

📄 論文解説: CachedAttention — マルチターン会話のKVキャッシュ再利用でTTFT 2.8倍高速化

本記事は CachedAttention: Cost-Efficient Large Language Model Serving for Multi-turn Conversations with KV Cache Reuse の解説記事です。

論文概要(Abstract)

CachedAttentionは、マルチターン会話においてセッション間のKVキャッシュを保存・再利用するシステムである。著者らは、従来の推論エンジンがリクエスト完了後にKVキャッシュを破棄する非効率に着目し、会話履歴のKV状態をCPUメモリやディスクにオフロードして次ターンで再利用する手法を提案している。vLLMをベースにした実装で、TTFT最大2.8倍高速化・スループット向上が報告されている。

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

情報源

背景と動機(Background & Motivation)

マルチターン会話はLLMアプリケーションの主要な利用形態である。ChatGPT、Claude、社内チャットボットなど、ユーザーは複数ターンにわたってLLMと対話する。各ターンでは、全会話履歴をプロンプトに含めて送信する必要がある。

従来の推論エンジン(vLLM、TGI等)には以下の課題がある:

  • KVキャッシュの揮発性: リクエスト処理完了後、生成されたKV状態はGPUメモリから解放される。次のターンでは全会話履歴のKV計算が再度行われる
  • プレフィックスキャッシュの限界: vLLMのプレフィックスキャッシュはGPUメモリ内でのみ動作するため、GPUメモリ容量に制約される。長い会話や同時接続数が多い場合、キャッシュが溢れて効果が低下する
  • TTFT(Time to First Token)の悪化: 会話が長くなるほどプレフィックスの再計算に時間がかかり、ユーザー体験が悪化する

著者らはこれらの課題に対し、GPUメモリ外(CPU/ディスク)へのKVキャッシュ永続化を提案している。

主要な貢献(Key Contributions)

  • 貢献1: マルチターン会話のKVキャッシュをCPU/ディスクに永続化し、次ターンで再利用するCachedAttentionシステムの設計と実装
  • 貢献2: KVキャッシュの効率的なオフロード/リロード機構(非同期転送、選択的キャッシュ)の提案
  • 貢献3: vLLMベースの実装でTTFT 2.8倍高速化を実証

技術的詳細(Technical Details)

マルチターン会話のKVキャッシュ問題

$t$ ターン目の会話では、全履歴 $H_t = (m_1, r_1, m_2, r_2, \ldots, m_t)$ をプロンプトに含める。ここで $m_i$ はユーザーメッセージ、$r_i$ はアシスタント応答である。

$t$ ターン目のKV計算コストは:

\[\text{Cost}(t) = O\left(\sum_{i=1}^{t-1} (|m_i| + |r_i|) + |m_t|\right) \cdot L \cdot d\]

ここで、

  • $m_i$, $r_i$: メッセージ/応答のトークン数
  • $L$: Transformerレイヤー数
  • $d$: 隠れ層次元数

会話が進むにつれて $\text{Cost}(t)$ は線形に増加する。例えば10ターンの会話で各ターン1,000トークンの場合、10ターン目では9,000トークン分のKV再計算が必要になる。

CachedAttentionのアーキテクチャ

CachedAttentionは以下の3つのコンポーネントで構成される:

1. KV Cache Manager(キャッシュマネージャ)

セッションIDをキーとして、各セッションのKVキャッシュ状態を管理する。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
from dataclasses import dataclass, field
from typing import Optional
import torch
import hashlib

@dataclass
class SessionKVState:
    """セッション単位のKVキャッシュ状態"""
    session_id: str
    turn_count: int
    total_tokens: int
    kv_tensors: Optional[dict[int, tuple[torch.Tensor, torch.Tensor]]] = None
    # layer_idx -> (K_tensor, V_tensor)
    storage_tier: str = "gpu"  # "gpu", "cpu", "disk"
    last_access_time: float = 0.0

class KVCacheManager:
    """マルチターン会話のKVキャッシュを管理"""

    def __init__(
        self,
        gpu_budget_tokens: int = 100_000,
        cpu_budget_tokens: int = 1_000_000,
    ):
        self.sessions: dict[str, SessionKVState] = {}
        self.gpu_budget = gpu_budget_tokens
        self.cpu_budget = cpu_budget_tokens

    def save_kv(
        self,
        session_id: str,
        layer_kvs: dict[int, tuple[torch.Tensor, torch.Tensor]],
        total_tokens: int,
    ) -> None:
        """リクエスト完了後にKV状態を保存

        通常のvLLMではここでKVが破棄されるが、
        CachedAttentionではCPUメモリに退避する
        """
        cpu_kvs = {
            layer: (k.cpu(), v.cpu())
            for layer, (k, v) in layer_kvs.items()
        }
        self.sessions[session_id] = SessionKVState(
            session_id=session_id,
            turn_count=self.sessions.get(session_id, SessionKVState(session_id, 0, 0)).turn_count + 1,
            total_tokens=total_tokens,
            kv_tensors=cpu_kvs,
            storage_tier="cpu",
        )

    def load_kv(
        self,
        session_id: str,
        device: torch.device,
    ) -> Optional[dict[int, tuple[torch.Tensor, torch.Tensor]]]:
        """次ターン処理時にKV状態をGPUにリロード"""
        state = self.sessions.get(session_id)
        if state is None or state.kv_tensors is None:
            return None

        gpu_kvs = {
            layer: (k.to(device), v.to(device))
            for layer, (k, v) in state.kv_tensors.items()
        }
        return gpu_kvs

2. Async Transfer Engine(非同期転送エンジン)

KVキャッシュのGPU→CPU転送をリクエスト処理と並列に行う。著者らは、CUDAストリームを活用した非同期メモリコピーにより、転送レイテンシをリクエスト処理時間に隠蔽する手法を提案している。

\[T_{\text{total}} = \max(T_{\text{inference}}, T_{\text{transfer}})\]

推論時間 $T_{\text{inference}}$ が転送時間 $T_{\text{transfer}}$ を上回る場合、KVオフロードのオーバーヘッドは実質ゼロになる。

3. Selective Cache Policy(選択的キャッシュポリシー)

全セッションのKVを保持すると膨大なメモリが必要になるため、著者らはLRU(Least Recently Used)ベースの選択的キャッシュポリシーを採用している。さらに、セッションの「継続確率」を推定し、次ターンが来る可能性が高いセッションのKVを優先的にCPUメモリに保持する。

推論フローの変更

従来のvLLMの推論フローとCachedAttentionのフローを比較する:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
【従来のvLLM】
ターン t:
  1. 全会話履歴 H_t をトークン化
  2. 全トークンのKV計算 → GPU使用
  3. 出力生成
  4. KVキャッシュを破棄 ← ここが非効率

ターン t+1:
  1. H_{t+1} をトークン化
  2. 再度全トークンのKV計算 ← H_t の分を重複計算
  3. 出力生成

【CachedAttention】
ターン t:
  1. 全会話履歴 H_t をトークン化
  2. 全トークンのKV計算 → GPU使用
  3. 出力生成
  4. KVキャッシュをCPUにオフロード(非同期) ← 保存

ターン t+1:
  1. H_{t+1} をトークン化
  2. CPUからKV_{1:t}をGPUにリロード ← 高速
  3. 新規トークン(m_{t+1}のみ)のKV計算 ← 差分のみ
  4. 出力生成

実装のポイント(Implementation)

  1. PCIe帯域幅のボトルネック: GPU→CPU間のKV転送はPCIe帯域(通常16-32 GB/s)に制約される。長い会話(100K+トークン)では転送に数百ミリ秒かかるため、非同期転送が必須

  2. KV圧縮との併用: FP16のKVキャッシュは1トークンあたり約128バイト(32層×128次元×2(K,V)×2バイト)。FP8やINT4量子化を適用することで転送量と保存容量を半減〜1/4に削減可能

  3. セッション管理の複雑さ: ステートレスなAPIサーバー(Lambda、ECS等)でセッション単位のKVキャッシュを管理するには、外部ストレージ(Redis、S3等)が必要。著者らはインメモリキャッシュを前提としており、分散環境での適用には追加の設計が必要

  4. vLLMとの統合: CachedAttentionはvLLMのPagedAttention機構と互換性があり、ページテーブルを拡張してCPU/ディスク上のKVブロックへの参照を管理する

実験結果(Results)

著者らはShareGPT会話データセットを用いた実験結果を報告している。

表: CachedAttentionの性能評価(論文より)

指標ベースライン(vLLM)CachedAttention改善率
TTFT(平均)ベースライン-最大2.8倍高速化
TTFT(P95)ベースライン-最大2.5倍高速化
スループットベースライン-向上(具体値は条件依存)

制約と注意点:

  • キャッシュのリロード時間はPCIe帯域に依存し、超長文会話ではオーバーヘッドが増加する
  • 同時接続セッション数が多い場合、CPUメモリが不足する可能性がある
  • 実験はオープンソースモデル(Llama系)で行われており、APIベースのClaude等では直接適用できない

実運用への応用(Practical Applications)

CachedAttentionの知見は、APIベースのLLMアプリケーション(Claude等)でも以下の形で活用可能である:

Anthropicプロンプトキャッシュとの関係: CachedAttentionはサーバー側のKV永続化であり、AnthropicのプロンプトキャッシュはAPI側でこの機能を提供している。ユーザー視点では、Anthropicのcache_controlパラメータを適切に設定することで、CachedAttentionと同等の効果を得られる。

LangGraphエージェントへの示唆:

  • マルチターン会話では、MemorySaverチェックポイントとcache_controlを組み合わせることで、会話状態の永続化とKVキャッシュの再利用を同時に実現できる
  • 5分TTLのキャッシュが切れないよう、長い処理(RAG検索等)中もキャッシュをアクティブに保つ設計が重要

スケーリング考慮:

  • 同時1,000セッション、各セッション平均5,000トークンの場合、KVキャッシュ総量は約640MB(FP16)。CPUメモリに収まる規模
  • セッション数が10万を超える場合、Redisやmemcachedを用いた分散KVキャッシュが必要

関連研究(Related Work)

  • vLLM Prefix Caching: GPUメモリ内でのプレフィックスキャッシュ。CachedAttentionはCPU/ディスクまで拡張した点が異なる
  • MemServe (2408.02357): 分散推論クラスタ間でKVキャッシュを弾力的に共有するフレームワーク。CachedAttentionの分散版と位置づけられる
  • RAGCache (2404.14294): チャンク単位のKVキャッシュ管理。CachedAttentionはセッション単位のキャッシュ管理に特化

まとめと今後の展望

CachedAttentionは、マルチターン会話におけるKVキャッシュの非効率(毎ターンの全履歴再計算)を、CPUオフロードによる永続化で解決した研究である。TTFT 2.8倍高速化という実用的な改善を実証しており、会話型LLMアプリケーションの推論最適化において重要な貢献である。

今後の研究方向として、KV圧縮との併用による転送効率の向上、分散環境でのセッション管理、APIプロバイダ(Anthropic等)への統合が挙げられる。

参考文献

この投稿は CC BY 4.0 でライセンスされています。