キャッシュ戦略

キーバリュー（KV）ベクトルは、アテンションスコアの計算に使用されます。自己回帰モデルでは、モデルが1トークンずつ予測するため、KVスコアは毎回計算されます。各予測は直前のトークルに依存するため、モデルは毎回同じ計算を実行します。

KV キャッシュは、これらの計算結果を保存し、再計算せずに再利用できるようにします。効率的なキャッシュは、計算時間を短縮し応答速度を向上させるため、モデルのパフォーマンス最適化において不可欠です。キャッシュの動作に関する詳細な説明については、Cachingドキュメントを参照してください。

Transformersには、異なるキャッシュ機構を実装した複数のCacheクラスが用意されています。これらのCacheクラスの中には、メモリを節約するように最適化されたものもあれば、生成速度を最大化するように設計されたものもあります。以下の表を参照してキャッシュの種類を比較し、ご自身のユースケースに最適なキャッシュを選択してください。

キャッシュタイプ	スライディングレイヤーをサポート	オフロードに対応	torch.compile() をサポート	予想されるメモリ使用量
Dynamic Cache	はい	はい	いいえ	中
静的キャッシュ	はい	はい	はい	高
量子化キャッシュ	いいえ	いいえ	いいえ	低

このガイドでは、さまざまなCacheクラスを紹介し、生成にそれらを使用する方法を示します。

デフォルトのキャッシュ

DynamicCache は、すべてのモデルのデフォルトのキャッシュクラスです。生成が進行するにつれて、キーとバリューの数を増加させて保存できるよう、キャッシュサイズを動的に拡張できます。

スライディングウィンドウアテンション（Mistral、Gemma2 など）またはチャンク化アテンション（Llama4）を使用するモデルでは、これらのアテンションタイプを使用するレイヤーが最大サイズ（スライディングウィンドウサイズまたはチャンクサイズ）に達すると、キャッシュの増加が停止します。

use_cache=False を設定して、generate() でキャッシュを無効にします。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

model.generate(**inputs, do_sample=False, max_new_tokens=20, use_cache=False)

Cache クラスを先に初期化してから、モデルの past_key_values パラメータに渡すこともできます。これにより、より細粒度の制御や、コンテキストキャッシュのような高度な利用が可能になります。

ほとんどの場合、cache_implementation パラメータでキャッシュ戦略を定義する方が簡単です。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

past_key_values = DynamicCache(config=model.config)
out = model.generate(**inputs, do_sample=False, max_new_tokens=20, past_key_values=past_key_values)

固定サイズキャッシュ

デフォルトの DynamicCache は、キャッシュサイズが固定されていないため、ほとんどのジャストインタイム (JIT) 最適化を活用できません。JIT 最適化は、メモリ使用量の増加を代償にレイテンシを最小限に抑えることを可能にします。以下のすべてのキャッシュタイプは、torch.compile などの JIT 最適化と互換性があり、生成を高速化できます。

固定サイズのカッシュ（StaticCache）は、kvペア用に特定の最大カッシュサイズを事前割り当てします。カッシュの修正を行うことなく、最大カッシュサイズまで生成できます。ただし、キー/バリュー状態のサイズが固定（通常は大きい）であるため、生成中に多くのトークンがアテンションに参加すべきではないためマスクされます。この手法により、デコーディング段階をcompileすることは容易ですが、アテンション計算においてトークンの無駄が発生します。これはすべてにおいてトレードオフであり、長さがほぼ同じ複数のシーケンスで生成する場合は非常に有効ですが、例えば非常に長いシーケンス1つと短いシーケンスのみがあるような場合は、最適ではない可能性があります（固定カッシュサイズが大きくなるため、短いシーケンスでは多くの部分が無駄になります）。使用する場合、その影響を十分に理解してください。

DynamicCache については、スライディングウィンドウアテンション（Mistral、Gemma2 など）またはチャンク化アテンション（Llama4）を使用するモデルでは、これらのアテンションタイプを使用するレイヤーにおいて、指定された最大長が大きい場合でも、キャッシュはスライディングウィンドウサイズ/チャンクサイズを超えることはありません。

generate() で cache_implementation="static" を設定することで、StaticCache を有効にできます。これにより、greedy 戦略および sample 戦略のデコーディング段階の自動compilationも有効になります。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("Hello, my name is", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="static")
tokenizer.batch_decode(out, skip_special_tokens=True)[0]
"Hello, my name is [Your Name], and I am a [Your Profession] with [Number of Years] of"

生成テンソルをCPUに保持する

Neuron や TPU などのコンパイラバックエンドは、モデルを固定された計算グラフにトレースします。生成ループでは、各ステップで1トークンずつ大きくなるテンソル（出力シーケンス、attention_mask、position_ids）を保持します。アクセラレーター上でこれらのテンソルの形状が変更されるたびにコンパイラがグラフを再トレースするため、生成が遅くなります。

generate() は、各 forward 呼び出しの直前に、forward が消費するテンソルのみをモデルのデバイスへ移動します。出力は、入力のデバイスに合わせて移動されます。ループ内で増大するテンソルの管理をアクセラレーターから外すには、入力を CPU に配置してください。コンパイルされたグラフは安定したまま保たれ、出力が入力のデバイスに従うため、生成された出力も CPU に留まります。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", device_map="auto")

# leave the inputs on CPU instead of calling .to(model.device).
inputs = tokenizer("The French Bread Law states", return_tensors="pt")

# generate runs forward on the model device but returns the output on the input_ids device.
output = model.generate(**inputs, do_sample=False, max_new_tokens=20)
print(output.device)
cpu

キャッシュのオフロード

KVキャッシュはメモリのかなりの部分を占有し、長文脈の生成におけるボトルネックとなる可能性があります。メモリ効率に優れたキャッシュは、メモリ使用量の削減と引き換えに速度を犠牲にするアプローチに焦点を当てています。これは、大規模言語モデル（LLM）や、ハードウェアがメモリ制約下にある場合に特に重要です。

キャッシュのオフロードは、1つのレイヤーを除くモデルレイヤーのKVキャッシュをCPUへ移動することで、GPUメモリを節約します。モデルがforwardでレイヤーを順に処理する間、GPU上に保持されるのは現在のレイヤーのキャッシュのみです。次のレイヤーのキャッシュは非同期で事前取得され、アテンション計算が完了すると、現在のレイヤーのキャッシュはCPUへ返送されます。

GPUが小さく、メモリ不足（OOM）エラーが発生している場合は、オフローディングを検討することをお勧めします。

モデルや生成の設定（コンテキストサイズ、生成トークン数、ビーム数など）によっては、完全なオンデバイスキャッシュと比較して、生成スループットがわずかに低下する場合があります。これは、キー/バリュー状態を往復させるのに一定の処理が必要だからです。

オフロードは、DynamicCache と StaticCache の両方で利用可能です。動的バージョンでは cache_implementation="offloaded"、静的バージョンでは cache_implementation="offloaded_static" を設定することで有効化でき、GenerationConfig または generate() のいずれかに設定できます。また、offloading=True オプションを指定して独自の DynamicCache または StaticCache をインスタンス化し、generate またはモデルの forward にこのキャッシュを渡すこともできます（例: 動的キャッシュの場合、past_key_values=DynamicCache(config=model.config, offloading=True)）。

上記で言及した2つのCacheクラスは、直接インスタンス化する場合に、offload_only_non_slidingという追加オプションを持っています。この追加引数は、スライディングウィンドウ/チャンクアテンションを使用しているレイヤー（存在する場合）もオフロードするかどうかを決定します。これらのレイヤーは通常、もともと短いため、オフロードには速度上のペナルティが発生する可能性があるため、オフロードを避ける方が望ましい場合があります。デフォルトでは、このオプションはDynamicCacheではFalse、StaticCacheではTrueです。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, dtype=torch.float16, device_map="auto")
inputs = tokenizer("Fun fact: The shortest", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=23, cache_implementation="offloaded")
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
Fun fact: The shortest war in history was between Britain and Zanzibar on August 27, 1896.

以下の例では、メモリが不足した場合にオフロードされたキャッシュにフォールバックする方法を示しています。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from accelerate import Accelerator

def resilient_generate(model, *args, **kwargs):
    oom = False
    device = Accelerator().device
    torch_device_module = getattr(torch, device, torch.cuda)
    try:
        return model.generate(*args, **kwargs)
    except torch.OutOfMemoryError as e:
        print(e)
        print("retrying with cache_implementation='offloaded'")
        oom = True
    if oom:
        torch_device_module.empty_cache()
        kwargs["cache_implementation"] = "offloaded"
        return model.generate(*args, **kwargs)

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, dtype=torch.float16, device_map="auto")
prompt = ["okay "*1000 + "Fun fact: The most"]
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
beams = { "num_beams": 40, "num_return_sequences": 20, "max_new_tokens": 23, "early_stopping": True, }
out = resilient_generate(model, **inputs, **beams)
responses = tokenizer.batch_decode(out[:,-28:], skip_special_tokens=True)

量子化キャッシュ

QuantizedCache は、KV 値を低精度に量子化することでメモリ要件を削減します。QuantizedCache は現在、3 つの量子化バックエンドをサポートしています:

hqqは、int1、int2、int3、int4、int8のデータ型をサポートしています。
quantoはint2およびint4のデータ型をサポートしています。これがデフォルトの量子化バックエンドです。
fp8 は KV 値を 8 ビット浮動小数点（float8_e4m3fn）で保存します。追加の依存関係は不要であり、グループ単位ではなくテンソル単位で量子化を行うため、nbits、axis_key、axis_value、q_group_size および residual_length は拒否されます。

コンテキスト長が短く、キャッシュの量子化を有効にしなくても生成に十分な GPU メモリがある場合、キャッシュの量子化はレイテンシを悪化させる可能性があります。メモリ効率とレイテンシのバランスを取るよう努めてください。

GenerationConfig で cache_implementation="quantized" を設定することで、QuantizedCache を有効にします。量子化バックエンドおよび追加の量子化パラメータは、generate() に渡す cache_config ディクショナリで指定します。これらの追加パラメータについては、メモリ不足が発生している場合を除き、デフォルト値を使用してください。メモリ不足が発生している場合は、残差長を短縮することを検討してください。

hqqバックエンドでは、axis_keyパラメータとaxis_valueパラメータを1に設定することをお勧めします。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=20,
    cache_implementation="quantized",
    cache_config={
        "nbits": 4,
        "backend": "hqq",
        "axis_key": 1,
        "axis_value": 1
    }
)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. It's a great way to express myself and rel

quantoバックエンドでは、axis_keyパラメータとaxis_valueパラメータを0に設定することをお勧めします。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=20,
    cache_implementation="quantized",
    cache_config={
        "nbits": 4,
        "backend": "quanto",
        "axis_key": 0,
        "axis_value": 0
    }
)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. It's a great way to express myself and rel

fp8バックエンドは追加のパラメータを受け取らず、グループワイズのパラメータが指定されると例外を発生させます。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Instruct-2507")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B-Instruct-2507", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="quantized", cache_config={"backend": "fp8"})
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it is powerful and emotional. I also like classical music because it is elegant and sophisticated. I enjoy

エンコーダー-デコーダーキャッシュ

EncoderDecoderCacheは、エンコーダー・デコーダーモデル向けに設計されています。自己注意機構と相互注意機構の両方のキャッシュを管理し、過去のkvペアの保存と取得を確実にします。エンコーダーとデコーダーに対して、それぞれ異なるキャッシュタイプを個別に設定することも可能です。

このキャッシュタイプにはセットアップは不要です。上記で説明した2つのCacheをラップするシンプルなもので、モデルが直接独立して使用します。

モデル固有のキャッシュ

一部のモデルには、他のどのキャッシュクラスとも互換性のない、独自の過去kvペアや状態の保存方法があります。

Mamba などの Mamba モデルは、アテンション機構や kv 状態を持たないため、専用のキャッシュが必要です。そのため、上記の Cache クラスとは互換性がありません。

反復生成

キャッシュは、モデルとの往復対話がある反復生成の設定（チャットボットなど）でも機能します。通常の生成と同様に、キャッシュを用いた反復生成では、各ステップでコンテキスト全体を再計算することなく、モデルが継続的な対話を効率的に処理できます。

キャッシュを用いた反復的な生成では、まず空のキャッシュクラスを初期化し、その後で新しいプロンプトを入力します。対話履歴はチャットテンプレートで管理します。

以下の例では、Llama-2-7b-chat-hf を使用しています。別のチャット形式のモデルを使用している場合、apply_chat_template() によるメッセージの処理方法が異なる可能性があります。Jinja テンプレートの記述方法によっては、重要なトークンが切り捨てられることがあります。マルチモーダルなチャットモデルについては、画像やオーディオの処理方法については、iterative chatting with cache ガイドを参照してください。

例えば、一部のモデルは推論中に<think> ... </think>のような特殊トークンを使用します。再エンコード時にこれらのトークンが失われる可能性があり、インデックスの問題を引き起こすことがあります。安定性を保つために、補完から余分なトークンを手動で削除または調整する必要がある場合があります。

import torch
from transformers import AutoTokenizer,AutoModelForCausalLM, DynamicCache, StaticCache

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(model_id)

user_prompts = ["Hello, what's your name?", "Btw, yesterday I was on a rock concert."]

past_key_values = DynamicCache(config=model.config)

messages = []
for prompt in user_prompts:
    messages.append({"role": "user", "content": prompt})
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
    input_length = inputs["input_ids"].shape[1]
    outputs = model.generate(**inputs, do_sample=False, max_new_tokens=256, past_key_values=past_key_values)
    completion = tokenizer.decode(outputs[0, input_length: ], skip_special_tokens=True)
    messages.append({"role": "assistant", "content": completion})

キャッシュをクロップする

crop() を使用して、キャッシュからトークンを巻き戻します。反復的な生成中に候補の継続が拒否された場合などに使用します。削除するトークン数を負の整数として渡します。

# Remove the last three tokens from the cache.
past_key_values.crop(-3)

Sliding-window 層と linear-attention 層は処理の進行に伴い過去の状態を破棄するため、デフォルトではロールバック先が存在しません。これらの層が状態を保持するようにするには、forward pass の前に activate_past_recording() を呼び出してください。これを呼び出さない場合、sliding-window 層がウィンドウを埋めると crop は RuntimeError を送出し、linear-attention 層では即座に送出されます。

# Enable this before the forward pass if you need to roll back tokens.
past_key_values.activate_past_recording()

crop(0) はキャッシュをクリアしません。キャッシュされたシーケンスからトークンを削除せず、フルアテンション層に対しては何も行わない操作（no-op）です。スライディングウィンドウ層や線形アテンション層では、次のモデル呼び出しで使用されないキャッシュ情報を破棄できます。これにより、トークンを巻き戻すことなくメモリ使用量を削減できます。

                    cached positions
                    ┌───┬───┬───┬───┬───┬───┐
before crop(-3)     │ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │
                    └───┴───┴───┴───┴───┴───┘
                                ╰─────┬─────╯
                                 rolled back

after crop(-3)
                    ┌───┬───┬───┐
full attention      │ 0 │ 1 │ 2 │    keeps the whole prefix
                    └───┴───┴───┘
                            ┌───┐
sliding_window=2            │ 2 │    keeps sliding_window - 1
                            └───┘

両方のレイヤーは現在、シーケンス長が3であることを報告し、生成は位置3から再開されます。

0 を渡すと、シーケンスからトークンは削除されません。フルアテンション層では何もしない操作ですが、スライディングウィンドウ層や線形アテンション層では、次のフォワードパスで読み込まれない記録済みの状態を破棄するため、ロールバックなしでメモリを解放できます。

キャッシュを事前入力する（プレフィックスキャッシュ）

状況によっては、特定のプレフィンプromptのkvペアでCacheを埋め、それを再利用して異なるシーケンスを生成したい場合があります。

以下の例では、StaticCacheを初期化し、初期プロンプトをキャッシュします。これにより、プリフィル済みプロンプトから複数のシーケンスを生成できます。

import copy
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, StaticCache

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map={"": 0})
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Init StaticCache with big enough max-length (1024 tokens for the below example)
# You can also init a DynamicCache, if that suits you better
prompt_cache = StaticCache(config=model.config, max_cache_len=1024)

INITIAL_PROMPT = "You are a helpful assistant. "
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to(model.device.type)
# This is the common prompt cached, we need to run forward without grad to be able to copy
with torch.no_grad():
     prompt_cache = model(**inputs_initial_prompt, past_key_values = prompt_cache).past_key_values

prompts = ["Help me to write a blogpost about travelling.", "What is the capital of France?"]
responses = []
for prompt in prompts:
    new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to(model.device.type)
    past_key_values = copy.deepcopy(prompt_cache)
    outputs = model.generate(**new_inputs, past_key_values=past_key_values,max_new_tokens=20)
    response = tokenizer.batch_decode(outputs)[0]
    responses.append(response)

print(responses)

Update on GitHub