Transformers 生成文本怎样理解 past_key_values 缓存

Transformers 文本生成中的 past_key_values 是已经处理过的上下文在注意力层里的 K/V 缓存。use_cache=True 让支持该接口的模型返回可复用缓存；要在手动调用中真正使用它，还得把缓存传回下一次调用，并且只传尚未处理的 token。它不是自动记住所有历史请求的全局开关。

一个容易出现的场景是：已经打开 use_cache，却在循环里每次把整个前缀重新送入模型。如果没传回缓存，历史仍可能被重新计算；若又传了缓存，又重复送入缓存已包含的 token，则可能重复计算位置，甚至产生长度错误。排查时先看输入与缓存分别代表哪些位置，而不是先猜“缓存没有生效”。

官方地址：https://huggingface.co/docs/transformers/main/cache_explanation

为什么打开 use_cache 后仍可能重复计算

把一次模型调用拆成三个对象看：input_ids 表示这次要处理的 token，past_key_values 表示已经处理的历史，输出中的 past_key_values 表示可用于后续调用的缓存。use_cache 是行为开关，不是缓存内容本身。

第一次没有历史缓存，可以输入完整提示词。下一次已经持有它的缓存，就只输入新增部分。新增部分通常是一个 token，也可以是多个尚未处理的 token；“使用缓存”不等于任何模型输入都只能有一个位置。

高层 generate 通常负责维护生成循环和这些状态，普通文本生成优先使用它。只有自己管理 forward 循环、续接前缀或排查张量长度时，才需要直接处理这些字段。本文讨论的是自回归推理，不把推理缓存当成训练加速配置。

past_key_values 保存的不是历史 token 文本

在这里讨论的因果自注意力中，历史 token 的 K 和 V 可以复用。缓存按注意力层组织，保存的是张量，不是原始字符串、完整聊天记录或模型权重。当前查询 Q 仍要结合相关的历史 K/V 计算注意力；缓存并没有消除所有上下文计算。

因此，input_ids、缓存 K/V、模型输出 logits 是三种不同数据。input_ids 用来表达本次新增输入；K/V 为注意力提供历史状态；logits 是模型对词表的分数，不是缓存对象。生成程序自己保存的完整 token 序列，也不能直接替代 past_key_values。

当前 GPT-2 文档将该参数描述为 Cache 实例，没有传入缓存时可以建立默认 DynamicCache。旧教程常见的嵌套 tuple 不能不看版本就照搬；也不要依赖某个模型内部层的具体索引布局来写通用代码。判断类型与长度应使用当前版本公开接口。

只把新增 token 交给已经持有历史的模型

下面的原创例子使用 GPT2Config 构造随机权重小模型，词表只含 32 个编号，不需要 tokenizer 或预训练权重。它不是聊天模型，也不用于展示生成质量。选择固定的额外 token 7，只是为了让缓存路径与完整前缀路径处理完全相同的四个位置。

示例面向当前 GPT-2 的 Cache 接口，需安装相应 Transformers 和 PyTorch。GPT-2 文档入口：https://huggingface.co/docs/transformers/v5.19.0/en/model_doc/gpt2。其他模型的位置处理和支持参数可能不同，不能把此例视为所有架构的通用生成循环。

import torch
from transformers import GPT2Config, GPT2LMHeadModel

# 随机权重小模型只说明缓存接口，不下载权重或生成有意义的文本
torch.manual_seed(11)
config = GPT2Config(
    vocab_size=32, n_positions=32, n_embd=32, n_layer=2, n_head=2,
    bos_token_id=0, eos_token_id=1, attn_implementation="eager",
)
model = GPT2LMHeadModel(config).eval()
prefix = torch.tensor([[4, 9, 6]], dtype=torch.long)
extra = torch.tensor([[7]], dtype=torch.long)

with torch.inference_mode():
    # 首次处理三个 token，取得它们的逐层 K/V
    first = model(
        input_ids=prefix, attention_mask=torch.ones_like(prefix),
        use_cache=True, return_dict=True,
    )
    cache = first.past_key_values
    if cache is None:
        raise RuntimeError("模型没有返回缓存，请确认模型接口与 use_cache 设置")
    print("初始缓存长度：", cache.get_seq_length())

    # 缓存已有三个位置，新增输入只有一个，mask 覆盖全部四个位置
    whole = torch.cat([prefix, extra], dim=1)
    full_mask = torch.ones_like(whole)
    continued = model(
        input_ids=extra, attention_mask=full_mask,
        past_key_values=cache, use_cache=True, return_dict=True,
    )

    # 无缓存路径独立重算完整前缀，仅比较同一末位的 logits
    reference = model(
        input_ids=whole, attention_mask=full_mask,
        use_cache=False, return_dict=True,
    )
    matches = torch.allclose(
        continued.logits[:, -1, :], reference.logits[:, -1, :],
        atol=1e-5, rtol=1e-4,
    )
    print("末位 logits 在容差内一致：", matches)
    print("更新后缓存长度：", continued.past_key_values.get_seq_length())

这里首次缓存对应三个 token；继续调用的 input_ids 长度是 1，而 attention_mask 长度是 4。完整参考路径输入长度为 4，并明确关闭缓存。比较的是同一个最终位置的词表分数，不是把两次输出的全部序列维度强行对齐。

如果读者在兼容环境执行，期望初始缓存长度为 3，续接后为 4，末位 logits 在合理浮点容差内一致。本文未执行代码，也没有提供实测打印结果。随机权重、数值后端和精度会影响具体分数；一致性比较不是速度测试。eval 关闭训练态随机行为，inference_mode 避免构建推理过程的梯度图，两者用途不同。

缓存输入出问题时先检查这三个长度

对于本例，先在调用之前核对缓存历史长度、新增输入长度、完整 attention_mask 长度。不要在调用后读取一个已被更新的可变缓存，再误以为那是调用之前的历史长度。

带 padding 的批量输入不能把所有 mask 都设为 1，应保留原有 padding 信息。本文全是有效 token，才使用 ones_like。位置参数也要按目标模型文档处理；如果自己的模型要求显式位置信息，需与缓存已有位置保持一致，不要机械沿用无 padding 的 GPT-2 示例。

缓存对象可能在调用中被更新。本例的 cache 只用于一次续接，然后使用 continued.past_key_values 表示更新后的状态。要拿同一前缀去试多个互不相关的分支，不能把同一个可变对象反复当成未改变的快照；应按缓存类型提供的复制方式隔离，或重新计算前缀缓存。

缓存为什么会增加内存占用？

它用保存历史 K/V 的空间换取少做重复计算。普通动态缓存会随序列增长；并不是开启缓存就能降低所有资源开销。长上下文、不同缓存策略和模型架构的内存表现不相同，不能把某次理论解释包装成固定倍数的实测提速。

换了提示词还能接着用旧缓存吗？

不能把不匹配的缓存当成新提示词的历史。缓存对应具体模型计算和已处理前缀，前缀内容、位置、批次或相关模型状态改变后，必须重新判断是否可复用。普通用法下新请求建立自己的缓存，避免跨请求误混。

use_cache=False 会自动清空之前持有的对象吗？

不能把这个开关理解为释放引用或清空历史的指令。若要对照无缓存重算，就像上面的 reference 一样，不传 past_key_values，并显式关闭本次缓存输出。释放缓存需要由应用结束对相应对象的持有，具体内存回收行为由运行环境决定。

理解这套接口的关键，是明确“已经处理的历史”与“这次新增的输入”各自在哪里。use_cache 决定是否产出缓存，past_key_values 承接历史；二者配合正确，才谈得上增量生成。

官方缓存机制说明：https://huggingface.co/docs/transformers/main/cache_explanation。缓存策略参考：https://huggingface.co/docs/transformers/main/kv_cache。