12. 推理引擎与键值缓存
训练结束后，模型参数不再变化。推理要做的是：给定一段提示，一个词元一个词元地生成回答。
如果每次都把整段历史重新送进模型，计算会越来越慢。nanochat 的推理引擎用**键值缓存（KV Cache）**解决这个问题。
读推理引擎说明nanochat/engine.py第 1–21 行没有键值缓存会怎样
最简单的生成方法是：
输入：法国的首都是
模型输出：巴
输入：法国的首都是巴
模型输出：黎
输入：法国的首都是巴黎
模型输出：。每一步都把完整序列重新送入模型。对于长度为 T 的序列，注意力需要计算整段历史；生成 N 个词元时，历史部分被重复计算了很多次。
语言模型有一个特点：前面词元的 Key 和 Value 一旦算出来，就不会因为后面新增词元而改变。既然如此，可以把它们保存下来，下一步直接复用。
键值缓存保存什么
Transformer 的每一层、每个位置都会产生一组 K 和 V。推理引擎为每一层预分配缓存：
K 缓存形状：[层数, B, T, 键值头数, 每个头的维度]
V 缓存形状：[层数, B, T, 键值头数, 每个头的维度]nanochat/engine.py第 82–121 行缓存只保存 K 和 V，不保存 Q。Q 是当前词元的查询，生成完当前词元后就不需要复用。
缓存占用可以用下面的公式估算：
字节数 = 2 × 层数 × 键值头数 × 每个头的维度 × 序列长度 × 每元素字节数前面的 2 表示 K 和 V 两份。GQA 会让键值头数小于查询头数，因此缓存更小。nanochat 默认训练没有开启 GQA，但推理引擎和缓存结构支持它。
预填充与逐词元生成
推理分为两个阶段：
1. 预填充
把整段提示一次送进模型：
<|bos|> <|user_start|> 法国的首都是哪里？ <|user_end|> <|assistant_start|>模型计算所有位置的 K 和 V，写入缓存，并输出最后一个位置的 logits。这个 logits 用来选第一个回答词元。
预填充的计算量大，但可以并行处理整段提示。首字延迟主要来自这一阶段。
2. 逐词元生成
之后每一步只送入一个新词元：
上一步生成的词元 → 模型 → 读取历史 K/V，追加当前 K/V → 新的 logits → 下一个词元因为历史 K/V 已经在缓存里，模型不需要重新投影它们，只需要为当前词元计算一份新的 Q/K/V。
推理时序
图 2：Engine.generate 的时序。预填充只做一次，生成循环每次处理一个新词元。
读 Engine.generatenanochat/engine.py第 169–275 行代码中有几个值得注意的细节：
- 预填充使用批大小为 1 的缓存。
- 如果要一次生成多个样本，会把预填充缓存复制到更大的解码缓存。
cache_seqlens记录每个样本当前的位置。当前实现假设同批样本的位置一致。- 模型处理完最后一层注意力后，缓存位置才前进，避免同一层重复推进。
- 生成结束后，
<|assistant_end|>或<|bos|>会把当前行标记为完成。
为什么缓存不是“把整段文本存起来”
缓存保存的是每一层的 K 和 V 张量，不是原始文字，也不是最终隐藏状态。它依赖模型结构、层数、头数和精度。
如果你更换模型、改变层数或修改注意力实现，旧的缓存就不能复用。缓存只在同一次推理过程中有效，不会被保存到磁盘，也不会跨进程复用。
工具调用状态机
nanochat 的 Engine 不只是采样词元，还负责处理 Python 工具调用。
当模型生成 <|python_start|> 时，Engine 进入 Python 块状态，开始收集后续词元作为表达式。遇到 <|python_end|> 时：
- 把收集到的词元解码成字符串。
- 用受限的
use_calculator()执行表达式。 - 如果得到结果，把
<|output_start|> 结果 <|output_end|>强制加入下一轮输入。 - 继续生成。
nanochat/engine.py第 240–267 行工具执行有白名单和 3 秒超时。它只支持基本算术和字符串的 .count()，不是完整的 Python 解释器。这样设计是为了降低执行任意代码的风险。
强制加入的输出词元会带 mask=0，表示它们不是模型采样的结果。训练时这些位置也不会被当成模型预测目标。
采样和缓存的关系
采样发生在模型输出 logits 之后：
logits → 温度缩放 → top-k 过滤 → softmax → 抽样 → 新词元新词元被送进下一轮模型，同时写入缓存。如此循环，直到遇到停止词元或达到最大长度。
下一章会详细讲温度、top-k 和命令行对话。
亲手观察缓存效果
可以比较两个生成函数：
GPT.generate()：不使用缓存，每一步重新计算整段序列。Engine.generate()：先预填充，再使用缓存。
nanochat/gpt.py第 526–555 行在 CPU 上，小模型生成几十个词元就能看出差异。注意比较相同提示、相同随机种子和相同最大长度下的耗时和输出。开启缓存后，输出应当基本一致，但速度更快。
本章要记住的三件事
- 键值缓存保存每层历史的 K 和 V，避免在生成新词元时重复计算历史。
- 推理分为预填充和逐词元生成两阶段；首字延迟主要来自预填充。
- Engine 还负责工具调用状态机，工具输出由引擎强制注入，不作为模型训练目标。
下一章看怎样控制生成结果：采样与命令行对话。