缓存策略

键值（KV）向量用于计算注意力分数。对于自回归模型，由于模型每次只预测一个 token，因此 KV 分数是每次都计算的。每个预测都依赖于之前的 token，这意味着模型每次都会执行相同的计算。

KV 缓存会存储这些计算结果，以便在不重新计算的情况下重复使用。高效的缓存对于优化模型性能至关重要，因为它能减少计算时间并提高响应速度。有关缓存工作原理的详细说明，请参阅缓存文档。

Transformers 提供了多个 Cache 类，它们实现了不同的缓存机制。其中一些 Cache 类针对节省内存进行了优化，而另一些则旨在最大化生成速度。请参阅下表以比较不同的缓存类型，并据此为您的使用场景选择最合适的缓存。

缓存类型	支持滑动层	支持卸载	支持 torch.compile()	预期内存使用量
动态缓存	是的	是的	No	中等
静态缓存	是的	是的	是的	高
量化缓存	No	否	No	低

本指南将为您介绍不同的 Cache 类，并展示如何在生成过程中使用它们。

默认缓存

DynamicCache 是所有模型的默认缓存类。它允许缓存大小动态增长，以便随着生成过程的推进，存储越来越多的键和值。

请注意，对于使用滑动窗口注意力（Mistral、Gemma2 等）或分块注意力（Llama4）的模型，当使用这些注意力类型的层达到其最大尺寸（即滑动窗口或分块大小）时，缓存将停止增长。

通过在 generate() 中设置 use_cache=False 来禁用缓存。

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 对预先分配一个特定的最大缓存大小。你可以生成多达最大缓存大小的 token，而无需对其进行修改。然而，由于 key/value 状态具有固定（通常较大）的大小，在生成过程中，许多 token 实际上会被屏蔽，因为它们不应参与注意力计算。因此，这种技巧可以轻松地 compile 解码阶段，但会在注意力计算中造成 token 的浪费。凡事皆有取舍：如果你生成的多个序列长度大致相同，这种方案效果会非常好；但如果你有一个非常大的序列，随后只有较短的序列（因为固定缓存大小会很大，短序列会浪费大量空间），则可能不是最优选择。如果你打算使用它，请务必了解其影响！

对于 DynamicCache，请注意，对于使用滑动窗口注意力（Mistral、Gemma2 等）或分块注意力（Llama4）的模型，即使指定的最大长度更大，在使用这些类型注意力的层上，缓存的大小也永远不会超过滑动窗口/分块大小。

您可以通过在 generate() 中设置 cache_implementation="static" 来启用 StaticCache。这还将为贪心解码和采样解码策略自动开启解码阶段的 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 这样的编译器后端会将模型追踪为固定的计算图。生成循环会维护一些张量（输出序列、attention_mask、position_ids），这些张量在每一步都会增加一个 token。每当这些张量在加速器上的形状发生变化时，编译器就会重新追踪计算图，从而拖慢生成速度。

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）以及硬件内存受限的场景，这一点尤为重要。

通过将缓存卸载来节省 GPU 显存：除一个模型层外，其余层的 KV 缓存会被移至 CPU。在模型对各个层执行 forward 迭代期间，GPU 上仅保留当前层的缓存。系统会异步预取下一层的缓存，并在注意力计算完成后，将当前层的缓存送回 CPU。

如果您的 GPU 较小且遇到内存不足（OOM）错误，可以考虑使用 offloading。

与完全在设备端使用缓存相比，你可能会注意到生成吞吐量略有下降，具体取决于你的模型和生成配置（上下文大小、生成的 token 数量、beam 数量等）。这是因为在设备间来回移动 key/value 状态需要额外的计算开销。

Offloading 功能同时支持 DynamicCache 和 StaticCache。您可以通过在 GenerationConfig 或 generate() 中配置 cache_implementation="offloaded"（动态版本）或 cache_implementation="offloaded_static"（静态版本）来启用该功能。此外，您也可以使用 offloading=True 选项自行实例化 DynamicCache 或 StaticCache，并将该缓存传入 generate 或模型的 forward（例如，对于动态缓存，使用 past_key_values=DynamicCache(config=model.config, offloading=True)）。

请注意，上述提到的两个 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 目前支持三种量化后端：

hqq 支持 int1、int2、int3、int4 和 int8 数据类型。
quanto 支持 int2 和 int4 数据类型。这是默认的量化后端。
fp8 以 8 位浮点格式（float8_e4m3fn）存储 KV 值。它无需额外依赖，并且采用按张量（per-tensor）而非按组（per-group）的量化方式，因此会拒绝 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 专为编码器-解码器模型设计。它同时管理自注意力（self-attention）和交叉注意力（cross-attention）缓存，以确保之前 kv 对的存储与检索。可以分别为编码器和解码器设置不同的缓存类型。

此缓存类型无需任何配置。它是对上述两个 Cache 的简单封装，将由模型直接独立使用。

特定模型的缓存

某些模型具有独特的过去 kv 对或状态的存储方式，与其他任何缓存类都不兼容。

像 Mamba 这样的 Mamba 模型需要特定的缓存，因为该模型没有注意力机制或 kv 状态。因此，它们与上述 Cache 类不兼容。

迭代生成

缓存同样适用于存在与模型来回交互的迭代生成场景（例如聊天机器人）。与常规生成类似，使用缓存的迭代生成使模型能够高效地处理持续进行的对话，而无需在每一步都重新计算整个上下文。

对于使用缓存的迭代生成，首先初始化一个空的缓存类，然后即可输入新的提示词。使用 chat template 来跟踪对话历史。

以下示例演示了 Llama-2-7b-chat-hf。如果您使用的是其他聊天式模型，apply_chat_template() 可能会以不同的方式处理消息。具体取决于 Jinja 模板的编写方式，它可能会截断重要的 token。对于多模态聊天模型，请参阅 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() 从缓存中回滚 token，例如在迭代生成过程中候选续写被拒绝时。将需要移除的 token 数量作为负整数传入。

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

滑动窗口和线性注意力层在计算过程中会丢弃过去的状态，因此默认情况下没有可供回滚的状态。请在前向传播之前调用 activate_past_recording()，使这些层保留这些状态。否则，一旦滑动窗口层填满其窗口，crop 就会抛出 RuntimeError；而对于线性注意力层，则会立即抛出该错误。

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

crop(0) 不会清除缓存。它不会从缓存序列中移除任何 token，对于全注意力层而言是一个空操作。对于滑动窗口和线性注意力层，它可以丢弃下一次模型调用不会使用的缓存信息。这在不回滚 token 的情况下降低了内存占用。

                    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 不会从序列中移除任何 token。对于全注意力层，这是一个无操作；但对于滑动窗口和线性注意力层，它仍会丢弃下一次前向传播不会读取的已记录状态，从而在不回滚任何内容的情况下释放内存。

预填充缓存（前缀缓存）

在某些情况下，您可能希望使用特定前缀提示的 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