Home 論文解説: Prompt Cache — モジュラーAttention再利用による低レイテンシLLM推論
投稿
キャンセル

📄 論文解説: Prompt Cache — モジュラーAttention再利用による低レイテンシLLM推論

本記事は Prompt Cache: Modular Attention Reuse for Low-Latency Inference の解説記事です。

論文概要(Abstract)

Prompt Cacheは、LLM推論におけるTime-To-First-Token(TTFT)を削減するための手法である。プロンプト内に繰り返し現れる共通セグメント(システム指示、Few-shot例、ドキュメントコンテキスト等)のAttention状態(KV Cache)を事前に計算・保存し、異なるリクエスト間で再利用する。著者らはPrompt Markup Language(PML)を提案し、プロンプトをモジュラーな再利用単位に分割する枠組みを定義している。LLaMA-2(7B/13B)での評価において、TTFTを最大8倍短縮し、生成品質の劣化がないことが報告されている。

この記事は Zenn記事: プロンプトキャッシュ×セマンティックキャッシュ:LLM2層キャッシュ戦略 の深掘りです。

情報源

  • arXiv ID: 2311.04934
  • URL: https://arxiv.org/abs/2311.04934
  • 著者: In Gim, Guojun Chen, Seung-seob Lee, Nikhil Sarda, Anurag Khandelwal, Lin Zhong
  • 発表年: 2023
  • 分野: cs.LG, cs.CL

背景と動機(Background & Motivation)

LLMの推論コストにおいて、プロンプトの処理(prefill phase)はトークン数に比例した計算量を必要とする。特にRAGやチャットボットのように、システムプロンプト・ドキュメントコンテキスト・会話履歴がリクエストごとに大量に付加されるアプリケーションでは、同一のテキストセグメントが何度も再処理される非効率が生じる。

従来のKV Cache再利用手法は、先行リクエストと完全に同一のプレフィックスが存在する場合にのみ適用可能であった。しかし実際のアプリケーションでは、プロンプト内のセグメントの順序や組み合わせがリクエストごとに変化する。例えば、Few-shot例の選択が異なったり、RAGで取得するドキュメントの組み合わせが変わったりする。この場合、プレフィックス全体が一致しないためKVの再利用ができず、毎回全トークンの再計算が必要となっていた。

Prompt Cacheはこの問題を、プロンプトをモジュラーなセグメント単位で管理し、各セグメントを位置独立(position-independent)に扱うことで解決する。

主要な貢献(Key Contributions)

  • Prompt Markup Language(PML)の提案: プロンプトを再利用可能なモジュールに分割するためのマークアップ言語。各モジュールにキャッシュキーを付与し、セグメント単位でのKV Cache管理を可能にする
  • 位置独立なAttention状態の実現: Transformerの位置エンコーディングをセグメント境界でリセットすることで、セグメントの並び順に依存しないKV Cacheの再利用を実現する
  • 最大8倍のTTFT削減: LLaMA-2(7B/13B)での実験において、生成品質の劣化なしにTTFTを最大8倍短縮したと報告されている

技術的詳細(Technical Details)

Prompt Markup Language(PML)

PMLは、プロンプトを木構造のモジュールに分解するためのスキーマである。各モジュールは一意のIDを持ち、そのテキスト内容に基づいたハッシュがキャッシュキーとなる。

1
2
3
4
5
6
7
8
9
10
11
<module id="system_instruction">
  あなたはECサイトのカスタマーサポートAIです。
  以下のFAQデータベースを参照して回答してください。
</module>
<module id="faq_database">
  [FAQ項目1] 返品ポリシー: ...
  [FAQ項目2] 配送について: ...
</module>
<module id="user_query">
  返品ポリシーを教えてください
</module>

この構造により、system_instructionfaq_databaseのKV Cacheは事前計算・保存が可能となる。新しいリクエストではuser_query部分のみを処理すればよい。

位置エンコーディングの処理

標準的なTransformerでは、各トークンに絶対位置エンコーディング(またはRoPE等の相対位置エンコーディング)が付与される。KV Cacheを再利用するには、キャッシュ時と利用時で位置エンコーディングが一致する必要がある。

Prompt Cacheでは、各モジュールの先頭で位置インデックスをリセットする方式を採用している。モジュール内のトークン $t_i$ に対する位置エンコーディングは以下のように計算される:

\[\text{pos}(t_i) = i - \text{start}(m)\]

ここで、

  • $t_i$: モジュール $m$ 内の $i$ 番目のトークン
  • $\text{start}(m)$: モジュール $m$ の開始位置(グローバルインデックス)
  • $\text{pos}(t_i)$: トークン $t_i$ に付与される位置インデックス(モジュール内ローカル)

この設計により、モジュールのKV Cacheは挿入位置に依存しなくなる。著者らは、この位置リセットがAttention計算に与える影響を実験で検証し、生成品質に有意な差がないことを確認したと報告している。

Attention計算の変更

標準的なSelf-Attentionでは、全トークンのKeyとValueに対してAttentionを計算する:

\[\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V\]

Prompt Cacheでは、キャッシュ済みモジュールのKV ($K_{\text{cached}}, V_{\text{cached}}$) と新規トークンのKV ($K_{\text{new}}, V_{\text{new}}$) を結合してAttentionを計算する:

\[K = [K_{\text{cached}} ; K_{\text{new}}], \quad V = [V_{\text{cached}} ; V_{\text{new}}]\] \[\text{Attention}(Q_{\text{new}}, K, V) = \text{softmax}\left(\frac{Q_{\text{new}} K^T}{\sqrt{d_k}}\right) V\]

ここで、

  • $Q_{\text{new}}$: 新規トークンのQuery行列
  • $K_{\text{cached}}, V_{\text{cached}}$: キャッシュから読み込んだKey/Value行列
  • $K_{\text{new}}, V_{\text{new}}$: 新規トークンのKey/Value行列
  • $d_k$: Key次元数
  • $[; ]$: 系列方向の結合(concatenation)

キャッシュ済みKVの計算をスキップできるため、prefillフェーズの計算量はキャッシュミスしたトークン数に比例する。

アルゴリズム

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
from dataclasses import dataclass


@dataclass
class CachedModule:
    """キャッシュ済みモジュールのKV状態を保持する。"""
    module_id: str
    key_states: list  # shape: (num_layers, seq_len, num_heads, head_dim)
    value_states: list


class PromptCache:
    """Prompt CacheのKV Cache管理を行うクラス。"""

    def __init__(self, max_cache_size: int = 1024):
        self.cache: dict[str, CachedModule] = {}
        self.max_cache_size = max_cache_size

    def lookup(self, module_id: str, content_hash: str) -> CachedModule | None:
        """モジュールIDとコンテンツハッシュでキャッシュを検索する。"""
        key = f"{module_id}:{content_hash}"
        return self.cache.get(key)

    def store(self, module_id: str, content_hash: str, kv: CachedModule) -> None:
        """KV状態をキャッシュに保存する。"""
        key = f"{module_id}:{content_hash}"
        if len(self.cache) >= self.max_cache_size:
            oldest_key = next(iter(self.cache))
            del self.cache[oldest_key]
        self.cache[key] = kv

    def process_prompt(self, modules: list[dict], model) -> tuple:
        """PMLモジュールリストを処理し、KV Cacheを構築する。

        Args:
            modules: PMLモジュールのリスト
            model: LLMモデル

        Returns:
            結合されたKV Cache
        """
        combined_kv = []

        for mod in modules:
            cached = self.lookup(mod["id"], mod["content_hash"])
            if cached is not None:
                # キャッシュヒット: KVを再利用
                combined_kv.append((cached.key_states, cached.value_states))
            else:
                # キャッシュミス: KVを計算してキャッシュ
                kv = model.compute_kv(mod["tokens"], position_offset=0)
                self.store(
                    mod["id"],
                    mod["content_hash"],
                    CachedModule(mod["id"], kv[0], kv[1]),
                )
                combined_kv.append(kv)

        return self._concat_kv(combined_kv)

    def _concat_kv(self, kv_list: list[tuple]) -> tuple:
        """複数モジュールのKVを系列方向に結合する。"""
        # 実装省略: レイヤーごとにKey/Valueを結合
        pass

実装のポイント(Implementation)

メモリ管理: LLaMA-2 13Bの場合、1024トークンのKV Cacheは約400MBのGPUメモリを消費する。キャッシュ対象モジュールの選定とLRU(Least Recently Used)ポリシーによるエビクションが不可欠である。

位置エンコーディングの互換性: RoPE(Rotary Position Embedding)を使用するモデルでは、位置リセットの実装にAttentionカーネルの修正が必要となる。著者らはLLaMA-2のRoPE実装を変更し、モジュール境界でのリセットを実現している。

キャッシュキーの設計: モジュールIDとコンテンツハッシュの組み合わせがキャッシュキーとなる。トークナイザの出力が一致する必要があるため、同一テキストでもトークナイザのバージョンが異なるとキャッシュミスが発生する。

制約事項: Prompt Cacheはトークン単位の完全一致を前提とする。意味的に類似しているが表現が異なるプロンプトにはキャッシュが効かない。この制約が、Zenn記事で解説されている「セマンティックキャッシュとの相補性」の根拠となっている。

Production Deployment Guide

AWS実装パターン(コスト最適化重視)

Prompt Cacheをプロダクション環境で運用する場合のAWS構成を示す。KV Cacheの永続化とリクエストルーティングが設計上のポイントとなる。

規模月間リクエスト推奨構成月額コスト主要サービス
Small~3,000 (100/日)Serverless$50-150Lambda + Bedrock + DynamoDB
Medium~30,000 (1,000/日)Hybrid$300-800Lambda + ECS Fargate + ElastiCache
Large300,000+ (10,000/日)Container$2,000-5,000EKS + Karpenter + EC2 Spot

コスト試算の注意事項: 上記は2026年4月時点のAWS ap-northeast-1(東京)リージョン料金に基づく概算値です。実際のコストはトラフィックパターン、リージョン、バースト使用量により変動します。最新料金は AWS料金計算ツール で確認してください。

Small構成の詳細(月額$50-150):

  • Lambda: 1GB RAM, 30秒タイムアウト ($20/月)
  • Bedrock: Claude 3.5 Haiku, Prompt Caching有効 ($80/月)
  • DynamoDB: On-Demand、キャッシュメタデータ管理 ($10/月)
  • CloudWatch: 基本監視 ($5/月)

コスト削減テクニック:

  • Bedrock Prompt Caching有効化で入力トークンコスト30-90%削減
  • Spot Instances使用で推論サーバーコスト最大90%削減(EKS構成時)
  • ElastiCache Serverlessで低トラフィック時の自動スケールダウン

Terraformインフラコード

Small構成(Serverless): Lambda + Bedrock + DynamoDB

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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
module "vpc" {
  source  = "terraform-aws-modules/vpc/aws"
  version = "~> 5.0"

  name = "prompt-cache-vpc"
  cidr = "10.0.0.0/16"
  azs  = ["ap-northeast-1a", "ap-northeast-1c"]
  private_subnets = ["10.0.1.0/24", "10.0.2.0/24"]

  enable_nat_gateway   = false
  enable_dns_hostnames = true
}

resource "aws_iam_role" "lambda_bedrock" {
  name = "prompt-cache-lambda-role"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Action    = "sts:AssumeRole"
      Effect    = "Allow"
      Principal = { Service = "lambda.amazonaws.com" }
    }]
  })
}

resource "aws_iam_role_policy" "bedrock_invoke" {
  role = aws_iam_role.lambda_bedrock.id
  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect   = "Allow"
      Action   = ["bedrock:InvokeModel", "bedrock:InvokeModelWithResponseStream"]
      Resource = "arn:aws:bedrock:ap-northeast-1::foundation-model/anthropic.claude-3-5-haiku*"
    }]
  })
}

resource "aws_lambda_function" "prompt_cache_handler" {
  filename      = "lambda.zip"
  function_name = "prompt-cache-handler"
  role          = aws_iam_role.lambda_bedrock.arn
  handler       = "index.handler"
  runtime       = "python3.12"
  timeout       = 60
  memory_size   = 1024

  environment {
    variables = {
      BEDROCK_MODEL_ID    = "anthropic.claude-3-5-haiku-20241022-v1:0"
      DYNAMODB_TABLE      = aws_dynamodb_table.cache_metadata.name
      ENABLE_PROMPT_CACHE = "true"
    }
  }
}

resource "aws_dynamodb_table" "cache_metadata" {
  name         = "prompt-cache-metadata"
  billing_mode = "PAY_PER_REQUEST"
  hash_key     = "module_hash"

  attribute {
    name = "module_hash"
    type = "S"
  }

  ttl {
    attribute_name = "expire_at"
    enabled        = true
  }
}

resource "aws_cloudwatch_metric_alarm" "lambda_cost" {
  alarm_name          = "prompt-cache-cost-spike"
  comparison_operator = "GreaterThanThreshold"
  evaluation_periods  = 1
  metric_name         = "Duration"
  namespace           = "AWS/Lambda"
  period              = 3600
  statistic           = "Sum"
  threshold           = 100000
  alarm_description   = "Lambda実行時間異常(コスト急増の可能性)"
  dimensions = {
    FunctionName = aws_lambda_function.prompt_cache_handler.function_name
  }
}

運用・監視設定

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import boto3

cloudwatch = boto3.client('cloudwatch')

# Bedrock Prompt Cache ヒット率監視
cloudwatch.put_metric_alarm(
    AlarmName='bedrock-cache-miss-rate',
    ComparisonOperator='GreaterThanThreshold',
    EvaluationPeriods=2,
    MetricName='CacheMissRate',
    Namespace='Custom/PromptCache',
    Period=300,
    Statistic='Average',
    Threshold=0.8,  # キャッシュミス率80%超でアラート
    AlarmDescription='Prompt Cacheヒット率低下'
)

コスト最適化チェックリスト

  • Bedrock Prompt Caching有効化(入力コスト最大90%削減)
  • システムプロンプトを最小キャッシュ長以上に設計(Claude: 2048トークン以上)
  • Lambda メモリサイズ最適化(CloudWatch Insights分析)
  • DynamoDB On-Demand(低トラフィック時コスト最適)
  • CloudWatch アラームでコスト異常検知
  • Bedrock usageフィールドでcache_read_input_tokensを定常監視

実験結果(Results)

著者らはLLaMA-2(7B/13B)を用いた実験結果を報告している。

設定ベースライン TTFTPrompt Cache TTFT高速化率
システムプロンプト 2K tokens1.2秒0.3秒4倍
システムプロンプト 4K tokens2.1秒0.35秒6倍
システムプロンプト 8K tokens3.8秒0.45秒約8倍

(論文の実験結果に基づく。数値はLLaMA-2 13Bの場合)

分析ポイント:

  • キャッシュ対象のプロンプト長が長いほど効果が大きい。8Kトークンのシステムプロンプトでは約8倍の高速化が報告されている
  • 出力トークンの生成速度はキャッシュの有無に影響されない(prefillフェーズのみが対象)
  • 生成品質(perplexity, ROUGE-L等)はベースラインと統計的に有意な差がなかったと著者らは報告している

実運用への応用(Practical Applications)

Prompt Cacheの概念は、2026年4月時点で主要LLMプロバイダの標準機能として実装されている:

  • Anthropic Claude: cache_controlパラメータによる明示的キャッシュ。読み取りコスト90%割引
  • OpenAI GPT-4o以降: 1024トークン以上で自動キャッシュ。50-90%割引
  • Google Gemini 2.5: 暗黙的+明示的の2層構造。90%割引(Gemini 2.5以降)

Zenn記事で解説されているように、プロンプトキャッシュはLLM推論を毎回実行する点が重要な制約である。セマンティックキャッシュと組み合わせた2層キャッシュ構成により、「同一コンテキストの再処理コスト」と「LLM呼び出し回数」の両方を削減できる。

関連研究(Related Work)

  • PagedAttention (vLLM): KV CacheをOSの仮想メモリのようにページ単位で管理し、メモリ断片化を排除する手法(Kwon et al., 2023)。Prompt Cacheはこの基盤上にモジュラーな再利用層を追加するものと位置づけられる
  • CacheBlend: RAGで取得した複数ドキュメントのKV Cacheを個別に事前計算し、クエリ時にブレンド(融合)する手法(Yao et al., 2024)。Prompt Cacheのセグメント単位再利用をRAGに特化させた発展形
  • CacheGen: KV Cacheをテンソル量子化・エントロピー符号化で圧縮し、サーバー間で転送する手法(Liu et al., 2023)。Prompt Cacheのストレージ効率を改善する相補的アプローチ

まとめと今後の展望

Prompt Cacheは、LLMのprefillフェーズにおけるKV Cache再利用をモジュール単位で実現する手法である。著者らはPMLによるセグメント分割と位置独立な再利用を提案し、LLaMA-2での実験でTTFTの最大8倍短縮を報告している。

この手法はトークン単位の完全一致を前提とするため、意味的に類似するが表現が異なるクエリには対応できない。Zenn記事で解説されている「セマンティックキャッシュ」との併用が、この制約を補完する2層キャッシュアーキテクチャの基盤となっている。

参考文献


:::message この記事はAI(Claude Code)により自動生成されました。内容の正確性については原論文を確認してください。 :::

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