Prompt Caching

プレフィックスキャッシングとも呼ばれるプロンプトキャッシングは、大規模言語モデル（LLM）がプロンプトの先頭部分に対して計算した中間結果を保存し、まったく同じテキストで始まる後続リクエストで再利用します。多くのリクエストは、システムプロンプト、ツール定義、例のセット、長い文書、それまでの会話など、変化しない長いプレフィックスを共有し、その後に短い新規部分が続きます。キャッシュがヒットすると、サーバーは新しいトークンだけを処理するため、最初の出力トークンまでの時間が短縮されます。ホスト型APIでは、キャッシュされた入力の料金も通常料金の一部に抑えられます。

プロンプトキャッシングの仕組み#

トランスフォーマー型LLMは、2段階で回答を生成します。プレフィルでは、各層がプロンプトの各トークンについてキーと値のベクトルを計算します。デコードでは、新しい各トークンがアテンション機構を通じて、それらすべてのキーと値に注目します。1つのリクエスト内では、KVキャッシュにこれらのテンソルが保持されるため、生成トークンごとにプロンプトを再計算する必要はありません。プロンプトキャッシングでは、リクエストをまたいでこれらを保持します。これらのモデルのアテンションは因果的であり、トークンのキーと値はその前にあるトークンだけに依存するため、同じプレフィックスを持つ2つのプロンプトでは、そのプレフィックスのキーと値が同一になり、保存済みデータで再計算を代替できます。

次の例では、Hugging Face TransformersのGPT-2を使って再利用の方法を示します。409トークンのシステムプロンプトを一度だけプレフィルし、キャッシュしたプレフィックスに各質問のトークンだけを追加して2つの質問に回答します。その結果をすべて再計算した場合と比較します。

import copy

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

# A long system prompt shared by every request, followed by different user questions
system = "You answer questions about a camera store. " + "Returns are accepted within 30 days with a receipt. " * 40
questions = [" Question: Can I return a lens?", " Question: Do you ship abroad?"]
prefix_ids = tokenizer(system, return_tensors="pt").input_ids

with torch.no_grad():
    # Prefill the shared prefix once and keep its key-value cache
    prefix_cache = model(prefix_ids).past_key_values
    for question in questions:
        ids = tokenizer(question, return_tensors="pt").input_ids
        # Cache hit: only the question's tokens run through the model
        cached = model(ids, past_key_values=copy.deepcopy(prefix_cache)).logits[0, -1]
        # No cache: the whole prompt is computed from scratch
        full = model(torch.cat([prefix_ids, ids], dim=1)).logits[0, -1]
        same = torch.allclose(cached, full, atol=1e-3)
        print(f"Computed {ids.shape[1]} of {prefix_ids.shape[1] + ids.shape[1]} tokens, same next-token scores: {same}")

Computed 8 of 417 tokens, same next-token scores: TrueとComputed 7 of 416 tokens, same next-token scores: Trueが出力されます。各リクエストでは、プロンプトの2%未満の部分だけをモデルに通しながら、全計算した場合と同じ次トークンのスコアが得られます。

オープンソースのモデルサービングエンジンは、この処理を自動的に適用します。vLLMはKVキャッシュを固定サイズのブロックに保持します。これはページドアテンションのメモリレイアウトで、新しい各リクエストに対して最長一致プレフィックスのブロックを再利用し、GPUメモリが不足すると最も長く使われていないブロックを退避します。SGLangのRadixAttentionはキャッシュ済みプレフィックスをラジックス木に保持し、先頭部分を共有するリクエストがそのプレフィックスを再利用できるようにします。SGLangの論文では、他の最先端の推論エンジンと比べて最大6.4倍のスループットが報告されています。In Gimらが執筆し、MLSys 2024で発表されたPrompt Cacheの論文では、テンプレートや文書など、宣言されたプロンプトモジュールにまで完全一致プレフィックス以外の再利用を拡張し、初回トークンまでの時間をGPUで8倍、CPUで60倍短縮したと報告しています。

LLM APIにおけるプロンプトキャッシング#

ホスト型APIでは、同じ仕組みを料金付きで利用できます。OpenAIのプロンプトキャッシングは、対応モデルでデフォルトで有効です。一方、AnthropicのClaudeプロンプトキャッシングは、リクエストごとに1回、または最大4つのコンテンツブロックに対して、cache_controlフィールドで有効化します。Anthropicでは、キャッシュ読み取りの料金は基本入力料金の10分の1以下、デフォルトの5分キャッシュへの書き込みは基本料金の1.25倍で、ヒットするたびに5分の有効期間がリセットされます。どちらのプロバイダーも、完全に一致するプレフィックスだけを照合します。Anthropicはツール、システムプロンプト、メッセージの順に比較するため、ツール定義を編集すると、それ以降にキャッシュされたすべての内容が無効になります。OpenAIのガイドでは、安定した指示と参照資料を先に置き、ユーザーごとの詳細を最後に置くことを推奨しています。

共有するプレフィックスが長いワークロードほど、最大の効果が得られます。

• AIエージェント：各ステップでツール定義と会話履歴全体を再送信します。会話履歴は、前のステップのプロンプトに最新のアクションとその結果を加えたものです。
• チャットボット: 会話の各ターンでは、前のターンのプロンプトに内容が追加されるため、新しいメッセージまでの履歴がキャッシュされます。
• ドキュメントの質問応答: 1つの長い契約書、レポート、リポジトリに関する多くの質問では、処理済みのドキュメントを再利用します。また、検索パイプラインは同じ箇所が再び取得されたときに、その箇所を再利用します。
• Few-shotプロンプト: 数十個の固定例を一度だけ計算し、すべての入力で再利用します。

オープンボキャブラリーのビジョンモデルも、同様の方法でプロンプト計算を再利用します。Ultralytics YOLOEは、set_classes()の呼び出し時に、MobileCLIP テキストエンコーダーを使って「forklift」などのクラスプロンプトを一度だけエンコードし、生成された埋め込みをモデルに保存します。そのため、以降の予測ではテキストエンコーダーを完全にスキップできます。save_prompt_embeddings()はこれらをNPZファイルに書き込み、後のセッションでエンコーダーを再実行する代わりにload_prompt_embeddings()で読み込みます。

プロンプトキャッシングと関連用語#

• KVキャッシュ: 1回の生成におけるトークン間でキーと値を再利用します。プロンプトキャッシングでは、別々のリクエスト間で再利用します。
• セマンティックキャッシング: 応答全体を保存し、新しいクエリが以前のクエリと同じ意味の場合に応答を返して、モデルの実行を完全に省略します。プロンプトキャッシングでは、新しいプロンプト部分に対してモデルを実行し、新しい回答を生成します。
• プロンプト圧縮: プロンプト自体を短縮する手法で、リクエストごとに変わるテキストに適しています。プロンプトの固定部分をキャッシュし、可変部分を圧縮すると、両方を効果的に組み合わせられます。
• プロンプトエンジニアリング: プロンプトの内容を決定します。キャッシングによって変わるのは、サーバーがプロンプトのどの部分を再計算するかだけです。キャッシュがヒットした場合、浮動小数点の丸め誤差を除き、次のトークンのスコアは変わりません。