3  训练与推理的两副面孔

本章地图

你在这里：数据流的两头 —— 第 5 步（怎么训）和第 7 步（怎么生成）。 本章解决：同一个模型，为什么训练时一次前向算完所有位置、推理时却要一个词一个词来；KV 缓存省下了什么、又漏掉了什么。 前置：第 1、2 章（知道一次前向里发生了什么）。 读完能做：解释”生成 1000 个字为什么要 1000 次前向”；并判断一个缓存实现是否真等价于无缓存版本。

一个”会背诗”却”只会吐标点”的模型。

第 1、2 章我们一直在做同一件事：把 loss 降下来。到第 2 章结束， 那个 2 层、4 头、d_model = 128、共 42 万参数的字符级模型， 在验证集上把 loss 压到了 4.141——比”闭着眼睛按字频猜”的 5.67 好了 1.53。

按照前面两章的标准，这个模型是训好了的。

然后我们让它写诗。给定”春江潮水连海平”，贪心解码，生成 50 个字：

春江潮水海平。，。，。，

五十步里，四十五步是同一个 <unk>，剩下五步在「。」和「，」之间来回跳。

这个落差不是”模型还不够好”能解释的。同一套权重，teacher forcing 下 它能给真实的唐诗打出 4.076 的 loss——它认得诗。可它自己写的时候， 写出来的东西连”不通顺”都算不上，那是一个死循环。

这一章要回答的就是这个落差。答案不在权重里，而在训练和推理是两副面孔： 训练时它一次看到整句话的所有位置，推理时它一次只能看到一个位置。 我们一直以为这两副面孔是同一件事的两种算法，代价只是速度—— 这一章会证明，代价远不止速度。

注记本章的代码与数据

# 需要 PyTorch。本机在 /home/mpp/micromamba/envs/dlspace/bin/python
python notebooks/run_ch3_case.py        # 本章全部实测数字
python notebooks/kv_cache.py            # KV 缓存的实现与四种 bug

数据导出到 data/ch3_case.json。本章正文里出现的每一个数字都来自这两个脚本。

3.1 先把现象摆清楚

3.1.1 起点：一个已经训好的模型

我们不改模型，不改数据，不改训练超参数。检查点是第 1 章留下的那个： 全唐诗 4003 首、25.5 万字、字符表 1200（覆盖语料 87.6%）， 2 层 Transformer、4 个头、d_model = 128、上下文长度 64，训练 2000 步。

设备 cuda:1　参数量 425,984　上下文 64　词表 1200
验证集 loss = 4.141

作为参照，同一个分词器下”按字频瞎猜”的交叉熵是 5.67。也就是说， 这个模型确实从语料里学到了东西——1.53 个 nat 的差距不是白来的。

3.1.2 现象：三种解码方式，三种完全不同的结果

同一个模型，同一个提示词”春江潮水连海平”（7 个字），生成 50 个字。 差别只在怎么从 logits 里挑下一个字。

第一种：贪心（greedy）——每一步都取概率最大的那个字。

春江潮水海平。，。，。，

这个输出需要解释一下：decode 会把 <unk>（编号 1）丢掉不显示。 实际上这 57 个 token 里，45 个是 <unk>，另外 12 个是「。」和「，」交替。 把它显示全，长这样：

春江潮水海平 <unk>。，<unk>。，<unk>。，<unk>。，<unk>。，…

第二种：贪心 + 屏蔽 <unk>——不许它输出”不认识的字”。 这是很多人遇到垃圾输出时的第一反应：把坏 token 禁掉。

春江潮水海平。不見昔年年不見，不見春風吹。不見君不得，不見君家。
不見君不見，不見，不如何處？何時？不見君不見，不

<unk> 一个都没有了。但读一下这段话：「不見」出现了 8 次。 它从”吐标点”变成了”念咒”。

第三种：采样（sampling）——不取最大，按概率抽。温度 τ = 0.8：

春江潮水海平，狂唱山城秋。作城外由流風日裏，北樓殘月在晴雲。
葉春無一日來，攀折綠江中。清誰與一留人，一來少年春。

τ = 1.2：

春江潮水海平梅水盡山春浦。君城外由還見溪裏，北陌殘江在浪飛。
葉春無十夫江，攀折綠江柳。清誰與狂江三何處？少裁春月

这不押韵，也不成篇，但它像诗：有”山城秋”“北樓殘月”“攀折綠江柳”这样的意象组合， 重复的字在正常范围内。

3.1.3 用数字说”有多重复”

图 3.1　 七种解码策略的重复程度。条形长度是 distinct-2（越高越不重复），右侧标出出现至少三次的最长重复片段。贪心和「贪心+屏蔽 unk」都是红色，四种采样策略都是绿色——分界线不在模型，在「怎么选字」。

“看起来像不像诗”太主观。我们需要一个能算的指标。

distinct-n：所有长度为 n 的片段里，有多少是不重复的。 一篇正常的中文文本，distinct-2（相邻两字组合的丰富度）通常在 0.9 以上； 如果一段文本里有一半的二字组合是重复的，那它一定在念咒。

再加一个指标：最长重复 n-gram——出现了至少 3 次的、最长的那个片段有多长。 正常文本这个数很小（4 已经算多了）；退化的文本会很长。

把七种解码策略都跑一遍（提示词相同，温度类参数固定随机种子）：

解码策略	distinct-1	distinct-2	distinct-3	最长重复	<unk> 占比
贪心	0.060	0.102	0.146	8-gram	88%
贪心 + 屏蔽 unk	0.340	0.531	0.667	4-gram	0%
τ = 0.8	0.760	1.000	1.000	无	0%
τ = 1.2	0.840	1.000	1.000	无	0%
top-k = 10	0.720	0.980	1.000	无	0%
top-p = 0.9	0.780	1.000	1.000	无	0%
top-p = 0.9, τ = 1.2	0.760	1.000	1.000	无	0%

这张表在说什么

• 贪心：distinct-2 = 0.102。一百个相邻二字组合里，九十个是重复的。 最长重复片段有 8 个字——它不是在写诗，它是在循环播放一段 8 字的录音。
• 屏蔽 <unk> 治不了本：distinct-2 从 0.102 升到 0.531，翻了好几倍， 但离正常文本的 0.9 还差得远，而且最长重复从 8-gram 变成 4-gram—— 它还是在循环，只是换了个循环单元。
• 只要引入随机性，退化就消失了：四种采样策略的 distinct-2 全部 ≥ 0.98， 最长重复 n-gram 直接变成”无”。连最”保守”的 top-k = 10 都一样。

所以问题不在”模型不会写”，而在“argmax”这个动作本身。

注记停下来想一想 ①

上面那张表里，top-k = 10 的 distinct-2 是 0.980，没有达到 1.000。 而 τ = 0.8 达到了 1.000。

两者都在”引入随机性”，为什么 top-k = 10 仍然有重复片段？

（提示：数一数 top-k = 10 的输出里「妾」字出现了几次。）

3.1.4 三个必须排除的嫌疑

在往下走之前，我们要排除三个”看起来很像元凶”的东西。

嫌疑一：模型没训好。

不成立。同一个模型，teacher forcing 下对真实唐诗的 loss 是 4.076。 如果模型没学到东西，这个数应该接近 5.67。

嫌疑二：<unk> 这个 token 本身是坏的。

部分成立，但不是根因。看一眼语料就明白了：

语料里 <unk> 占 18.9%

字表只保留最高频的 1200 个字，覆盖 87.6% 的字符——剩下 12.4% 全都映射到 <unk>。 所以在训练数据里，<unk> 本身就是一个出现频率接近 19% 的”字”。 模型在某个位置把它选成 argmax，从统计上说完全合理：它确实是那个位置最可能的输出。

但屏蔽它之后，输出从”吐标点”变成了”念咒”——根因不在这里。

嫌疑三：模型进入了某个”坏状态”，是碰巧。

这是最需要认真对待的一个。如果只是碰巧，那换个提示词应该就好了。

警告一个反直觉的验证

把提示词换掉，用贪心重新生成，会得到完全不同的循环单元—— 但一定会有循环。

我们试过以”床前明月光”“白日依山尽”“两个黄鹂鸣翠柳”开头， 得到的循环单元分别是「不知」「山中日」「春風」这类二字或三字组合。 循环的内容变了，循环这件事没变。

这不是”碰巧陷入坏状态”，这是确定性解码在退化分布上的必然结果。 下一节开始解释为什么。

3.1.5 本章要回答的问题

现象摆完了，问题也就清楚了：

1. 为什么训练时一切正常，推理时却崩了？ 训练和推理到底哪里不一样？
2. 为什么是”循环”这么具体的形式？ 为什么不是”越写越乱”，而是”卡住不动”？
3. 为什么加一点随机性就好了？ 采样凭什么能打破循环？
4. 怎么让它跑得够快，快到能多试几次？ ——这一问引出 KV 缓存。

第 1 问和第 2 问在 3.2、3.3 节回答；第 3 问在 3.5 节； 第 4 问在 3.4 节，而且答案会出乎意料（剧透：省了 28 倍计算，一点没变快）。

3.1.6 完整的实验设置

本章所有实验都基于同一个检查点，参数如下。列出完整配置是因为 本章的每一个数字都依赖它——换一个模型，这些数字会变。

项	值
语料	全唐诗 4003 首，25.5 万字
分词	字符级，字表 1200（覆盖率 87.6%）
层数	2
注意力头	4（每头 32 维）
d_{model}	128
d_{ff}	256
上下文长度	64
位置编码	可学习的绝对位置嵌入（不是 RoPE）
参数量	425,984
训练步数	2000
学习率	3e-3，warmup 100 步，余弦衰减
验证集 loss	4.141

三个选择需要说明。

为什么用这么小的模型？ 因为本章要讲的是机制，不是性能。 42 万参数的模型在本机 GPU 上跑一次完整实验只要几十秒， 可以反复做对照。而本章的核心结论（贪心必然循环、 缓存是恒等变换、顺序陷阱）和模型规模无关—— 换成 GPT-2 结论一样，只是数字不同。

为什么用字符级而不是 BPE？ 因为字符级让”下一个 token”这件事 对人类可见。你能直接读出”模型想写哪个字”， 不需要查 tokenizer。这对理解机制很有帮助。

代价是 <unk> 问题被放大了。 字表只保留 1200 个高频字， 覆盖 87.6%——剩下 12.4% 全变成 <unk>。 真实的 BPE 分词器覆盖率通常在 99.9% 以上，不会有这么严重的 <unk> 问题。

这个”缺陷”恰好成了本章案例的引子（见 3.1.2 节）， 但它确实是一个本可以避免的人为因素。读到后面请记住这一点。

注记停下来想一想 ②

上面的设置里，“上下文长度 64”和”位置编码用可学习的绝对位置嵌入” 这两项共同导致了 3.9.6 节那个 CUDA 崩溃。

如果换成 RoPE，那个崩溃还会发生吗？

（答案：不会以同样的形式发生。

RoPE 不给每个位置存一个嵌入向量，而是按位置旋转 Q 和 K。 所以没有”第 65 行不存在”这个问题——位置 65 只是旋转角度不同而已， 计算上完全合法。

但会有别的问题：模型训练时只见过位置 0–63 的旋转角度， 位置 65 的旋转是外推。它的行为没有保证—— 可能还行，可能完全乱掉。

所以 RoPE 把”硬崩溃”换成了”静默劣化”。 回顾 3.6.3 节那张表：静默劣化比崩溃更难发现。

这是个反复出现的权衡：让错误早点暴露，还是让它晚点暴露？ 本书的立场是前者——宁可崩，不可悄悄错。）

3.1.7 补充排查：换提示词、换模型、换热身

3.1.4 节排除了三个嫌疑。还有三个必须排除的，否则”这是确定性解码的必然” 这个结论站不住。

排查四：换提示词。

如果循环是”这个提示词碰巧触发的”，那换个开头应该就好了。

实测：以「床前明月光」「白日依山尽」「两个黄鹂鸣翠柳」开头， 贪心输出分别陷入「不知」「山中日」「春風」的循环。 循环单元变了，循环这件事没变。

排查五：换温度以外的解码参数。

如果只是 top-k/top-p 的某个配置有问题，那不用它们应该就好了。

实测：不用任何截断，只用温度（τ = 0.8、τ = 1.2）， distinct-2 都是 1.000。问题不在截断，在”是否随机”。

排查六：换模型规模——这一条最难，但最有说服力。

如果这是”42 万参数的小模型太笨”，那大模型应该不会。 本章没有重训一个大模型（成本不允许），但有两类旁证：

旁证一：第 2 章的实验里，同一个模型在不同深度下 （6 / 12 / 24 层）的贪心输出都退化。深度不改变结论。

旁证二：文献。Holtzman 等人 2020 年在 GPT-2（15 亿参数） 上复现了同样的现象——贪心和 beam search 都退化， 只有采样能解决。规模差 3500 倍，结论一样。

提示为什么”换模型”这条必须认真对待

在排查中，“是不是模型太小了”是一个万能的借口。 它可以解释任何失败，因此什么都不解释。

排除它的方法有两个： 1. 换一个真正不同的模型（本章做不到，成本不允许） 2. 找到文献里的同现象（本章的做法）

第 2 种方法的价值在于：如果别人在 3500 倍大的模型上看到同样的事， 那就不是规模问题。 但它有个前提—— 你得真的去读那篇论文，而不是凭印象引用。

本章 3.11 节列出的是我们实际读过并核对过的文献。

六条排查的总结

#	嫌疑	结论	依据
1	模型没训好	排除	teacher forcing loss 4.076
2	<unk> 是坏的	部分成立，非根因	屏蔽后仍在循环
3	碰巧进入坏状态	排除	换提示词仍循环，只是换了单元
4	特定提示词触发	排除	同 3
5	截断参数配置错	排除	只用温度也是好的
6	模型太小	排除	第 2 章多深度 + 文献佐证

六条排查，每一条都对应一个具体的验证动作和一组具体数字。 这就是”排查”和”猜”的区别——猜不需要数字，排查需要。

3.1.8 怎么复现这一章的全部数字

本章出现的每一个数字都来自两个脚本。复现步骤：

cd /home/mpp/ai-mechanisms-book

# 全章实测数据（约 3 分钟，主要花在时序基准上）
/home/mpp/micromamba/envs/dlspace/bin/python notebooks/run_ch3_case.py

# 缓存实现与四种 bug 的最小示例（约 30 秒）
/home/mpp/micromamba/envs/dlspace/bin/python notebooks/kv_cache.py

输出：data/ch3_case.json（约 12 KB），包含：

键	内容	对应正文
decoding	七种解码策略的退化指标与文本	3.1.2 / 3.1.3
exposure	扰动实验与置信度对比	3.7.2 / 3.7.3
cache_exact	预填充与增量的逐位差异	3.4.3
cache_bugs	五种实现的 logits 差异	3.9.1
mask_noop	掩码在单步 vs 分块下的差异	3.9.4
cache_timing	batch 1/16/64 的墙钟	3.4.4
cache_scaling	两种模型规模的长度扫描	3.4.4
sampling	熵-温度曲线、核曲线、顺序陷阱	3.5.2 / 3.5.4 / 3.5.5

提示跑之前先看这三件事

一、设备。 脚本默认 cuda:1（本机的 3090）。 A6000（cuda:0）显存常被占 85%，不要用。 加 --cpu 可以强制 CPU——时序数字会完全不同， 但所有正确性数字（logits 差异、退化指标）一个都不变。

二、时间。 全量跑约 3 分钟。其中 cache_scaling 占了大头 （要建两个随机权重的模型并做长度扫描）。 如果只想验证正确性，把 main() 里那两段 cache_scaling 删掉， 30 秒就能跑完。

三、随机性。 所有采样都用了固定种子， 所以同一个 (提示词, 参数) 组合永远给同样的输出。 这是刻意的——本书里引用的每一段生成文本， 你都能原样复现出来。

如果复现不出来，最可能的原因是 PyTorch 版本不同导致 torch.multinomial 的随机数序列不同。这时请以 退化指标（distinct-n、最长重复） 为准—— 它们对具体抽到哪个字不敏感，而正文的结论都建立在这些指标上。

3.1.9 本章的路线图

本章的结构是”从现象到机制再到工程”，每一节都在回答 3.1.5 节提出的 四个问题之一。先给一张地图，读的时候可以随时回来看。

节	回答哪个问题	一句话结论
3.2 训练那一面	训练和推理哪里不一样	训练一次算完所有位置，用因果掩码防止作弊
3.3 推理那一面	为什么推理必须串行	要预测第 t+1 个，需要第 t 个；而第 t 个还没生成
3.4 KV 缓存	怎么让它跑得快	过去位置的 K/V 永不变（可证明），存下来即可
3.5 采样的数学	为什么加随机性就好了	循环来自 argmax 的确定性；采样打破它
3.6 上下文长度	硬边界在哪	位置嵌入表是硬墙，注意力的 O(T^2) 是软墙
3.7 训练与推理的错配	代价是什么	训练时看真实前文，推理时看自己造的前文
3.8 交互实验	—	三个可拖动的实验
3.9 失效边界	—	缓存写错的六种样子 + 排查手册
3.10 上手实现	—	一个 20 行的缓存 + 一个完整的生成器
3.11 延伸阅读	—	五篇论文，每篇一句话说清读它能得到什么

读这一章需要的前置知识

必须有： - 第 1 章的注意力计算（QK^\top、softmax、加权求和） - 第 2 章的 Block 结构（残差、LayerNorm、FFN） - 基本的矩阵形状感觉（哪一维是 batch、哪一维是时间）

最好有： - 交叉熵的定义（没有的话 3.2.1 节会补） - 概率论里的期望和条件分布记号

不需要： - 任何关于”推理优化”的先验知识。本章从零开始。 - 读过 Transformer 原论文。所有需要的地方都会重述。

如果第 1、2 章没读：3.2 和 3.3 节仍然能读—— 它们讲的是”同样一次前向，训练和推理的用法有什么不同”， 不依赖注意力的内部细节。但 3.4 节之后建议先补第 1 章。

提示三条阅读路径

路径 A（只想搞懂现象）：3.1 → 3.3 → 3.5 → 3.7。 大约 40% 的篇幅，回答”为什么会循环、为什么采样有用”。

路径 B（要做推理优化）：3.1 → 3.3 → 3.4 → 3.6 → 3.9 → 3.10。 跳过采样和错配，直奔缓存和它的失效边界。

路径 C（完整读）：从头到尾。3.7.7 节是全章的收束， 建议不要跳过——它把前面所有现象归结到一个根源上。

三条路径都要读 3.1。 没有案例的具体数字， 后面的推导会失去锚点。

3.2 训练那一面：一次算完所有位置

3.2.1 目标函数：把”预测下一个词”变成一个数

先把训练在做什么写清楚。

给定一段文本 x_1, x_2, \dots, x_T，语言模型的训练目标是： 在每一个位置 t，用前 t 个字符预测第 t+1 个字符。

模型在位置 t 输出的是一个长度为 V 的向量 \mathbf{z}_t \in \mathbb{R}^V， 叫 logits（逻辑值）。V 是词表大小，本模型 V = 1200。 经过 softmax，它变成”下一个字是第 v 个字”的概率：

p_\theta(x_{t+1} = v \mid x_1, \dots, x_t) = \frac{\exp(z_{t,v})}{\sum_{u=1}^{V} \exp(z_{t,u})} \tag{3.1}

所以这个式子在说：模型对每个可能的字给一个分数，softmax 把分数变成概率。 分数高的字概率大。分母那个求和是为了让所有概率加起来等于 1。

训练用的损失是交叉熵——真实下一个字的负对数概率，在所有位置上取平均：

\mathcal{L} = -\frac{1}{T}\sum_{t=1}^{T} \log p_\theta(x_{t+1} \mid x_1, \dots, x_t) \tag{3.2}

所以这个式子在说：如果模型给真实的下一个字的概率是 1，这一项就是 0（不罚）； 给的概率越小，罚得越狠，而且是对数尺度地狠——给 0.01 的惩罚是给 0.5 的两倍多。

用一个具体数字走一遍

假设某个位置，模型输出的 logits 是 [2.0, 1.0, 0.0]，真实答案是第 1 个（下标 0）。

三个指数：e^{2.0} = 7.389，e^{1.0} = 2.718，e^{0.0} = 1.000。 分母 = 7.389 + 2.718 + 1.000 = 11.107。

概率：p_0 = 7.389/11.107 = 0.665，p_1 = 0.245，p_2 = 0.090。

交叉熵 = -\log(0.665) = 0.408。

如果模型更自信一点，logits 是 [4.0, 1.0, 0.0]： 分母 = 54.598 + 2.718 + 1.000 = 58.316，p_0 = 0.936，loss = 0.066。 “对的字概率更高”这件事，在对数尺度上被放大了。

反过来，如果模型押错了宝，logits 是 [0.0, 1.0, 2.0]（把第 3 个当成了答案）： p_0 = 0.090，loss = -\log(0.090) = 2.41。罚了 6 倍。

3.2.2 把 teacher forcing 跑一遍：逐位置的实况

抽象讲完了，看真实数据。

从验证集里取一段 10 个字的真实文本：「已驚飛。方從列子御」 （第 4 个字是 <unk>，也就是”字表里没有这个字”）。

模型对每个位置的预测如下。第二列是”到这一位为止的上文”， 第三列是真实答案，第四列是交叉熵损失，最后是模型给出的前三个候选：

位置	上文	真值	损失	模型的前三候选（概率）
0	已	驚	6.379	「，」0.142　<unk> 0.097　「。」0.028
1	已驚	飛	5.184	「。」0.274　<unk> 0.101　「，」0.085
2	已驚飛	。	1.676	「。」0.187　<unk> 0.107　「，」0.038
3	已驚飛。	<unk>	0.000	<unk> 1.000　「不」0.000　「日」0.000
4	已驚飛。	方	6.013	<unk> 0.145　「日」0.017　「不」0.009
5	已驚飛。方	從	4.533	<unk> 0.182　「知」0.102　「聞」0.029
6	已驚飛。方從	列	6.265	<unk> 0.163　「萬」0.016　「有」0.016
7	已驚飛。方從列	子	8.184	<unk> 0.175　「玉」0.038　「羽」0.033
8	已驚飛。方從列子	御	6.984	<unk> 0.257　「，」0.064　「重」0.016

整段平均损失 5.024，困惑度 152.07。

这张表里有三件事，每一件都比它看起来更重要

第一，位置 3 的损失是 0.000476——模型给出了 1.000 的概率。

它完全确定下一个字是 <unk>。而且它确实”对了”， 因为真实答案就是 <unk>。

但请想一想这意味着什么：模型学会了一件完全正确、却毫无用处的事—— “当一个位置的字不在字表里时，就输出 <unk>”。 语料里 18.9% 的位置都是这种情况，所以这个策略的期望收益非常高。

这就是 3.1 节贪心解码输出 88% <unk> 的直接原因。 不是模型坏了，是它在忠实地执行训练数据教给它的最优策略。

第二，位置 7 的损失是 8.184——模型认为「子」的概率不到千分之一。

e^{-8.184} \approx 2.8 \times 10^{-4}。而它给 <unk> 的概率是 0.175， 差了 600 倍。

「方從列子御」是一句正常的诗（列子御风，出自《庄子》）。 模型不是不懂这句话——它在位置 5 给「知」0.102、位置 1 给「。」0.274， 说明它对诗的语感是有的。它是在”具体是哪个字”上被字表卡住了。

第三，损失在 0.000 到 8.184 之间横跨了四个数量级。

平均值 5.024 是一个极具误导性的汇总。 如果用困惑度 152 来描述这个模型，你会以为它”在所有位置都很差”； 实际上它在标点位置（位置 2，损失 1.676）相当准， 在罕见字位置（位置 7，损失 8.184）完全抓瞎。

这就是为什么本书反复强调要看分布，不要只看均值。

注记停下来想一想 ③

位置 3 的损失是 0.000476，看起来几乎是 0。为什么不是恰好 0？

（提示：模型给 <unk> 的概率是 1.000，但那是显示到三位小数的结果。 真实概率是 1 - \epsilon，其中 \epsilon 是个很小的正数—— softmax 永远不会输出精确的 1，因为它要除以一个包含所有 1200 项的和。

-\log(1 - \epsilon) \approx \epsilon，所以损失 0.000476 意味着 真实概率约 1 - 4.76 \times 10^{-4}。

这是浮点数和 softmax 的性质，不是模型”不确定”。

3.2.3 关键问题：为什么训练能一次算完

现在到了本章的核心对比。

式 式 3.2 里有一个求和，t 从 1 到 T。最直接的做法是写个循环：

loss = 0
for t in range(T):                    # T 次前向
    logits = model(x[:t+1])           # 只看前 t+1 个字符
    loss += -log_softmax(logits[-1])[x[t+1]]
loss /= T

这是推理的写法（下一节会讲）。训练时的写法完全不同：

logits, loss = model(x)               # 1 次前向，T 个位置同时出结果

一次前向，输入是完整的 (B, T)，输出的 logits 形状是 (B, T, V)—— 每一个位置都有自己的预测。然后一次性算交叉熵。

为什么可以这样？因为 Transformer 在位置 t 的输出，只依赖于位置 1 \dots t 的输入。 只要保证这一点，把整句话一起喂进去，第 t 个位置的输出和”只喂前 t 个字符” 得到的输出就是同一个东西。

保证这一点的机制就是因果掩码（causal mask）。

掩码在做什么

注意力那一层，位置 t 可以”看”所有位置。如果不加限制， 它在算第 1 个字的预测时就会偷看第 2、3、4 个字——作弊， 而且推理时根本拿不到未来的字。

因果掩码的做法：在 softmax 之前，把所有”未来位置”的打分设成 -\infty。 softmax 之后 e^{-\infty} = 0，那些位置的权重就真的变成 0 了。

用矩阵看，一个 T = 4 的掩码长这样（1 = 允许看，0 = 屏蔽）：

M = \begin{pmatrix} 1 & 0 & 0 & 0 \\ 1 & 1 & 0 & 0 \\ 1 & 1 & 1 & 0 \\ 1 & 1 & 1 & 1 \end{pmatrix} \tag{3.3}

第 1 行只能看第 1 列，第 2 行能看前 2 列……第 4 行能看全部。 每一行恰好对应”从这个位置往回看，历史有多长”。

这个矩阵是下三角的，而且对角线全是 1。它的形状本身就说明了因果性： 位置越靠后，能看到的历史越长；没有任何一行能看到自己右边的位置。

注记停下来想一想 ④

式 式 3.3 的矩阵里，第 4 行全是 ✓。这看起来像”第 4 个位置偷看了未来”—— 毕竟它能看到全部 4 个位置。

它作弊了吗？

（提示：第 4 个位置预测的是第 5 个字，不是第 4 个字。它能看到的位置 是 1、2、3、4——全都是它的”过去”。）

3.2.4 一次前向，T 个损失，为什么不重复计算

真正让并行训练成为可能的，是一个更朴素的事实：

位置 t 的表示，和位置 t' 的表示，在前向传播中是各自独立算出来的。

自注意力把 T 个位置一次性投影成 Q, K, V 三个矩阵，然后是一次矩阵乘法 QK^\top。前馈网络是逐位置的，T 个位置各自过一遍同样的 MLP。 没有任何一步需要”先算完位置 1 再算位置 2”。

于是整个训练前向的计算量是：

• 注意力：O(T^2 \cdot d)（每对位置算一次点积）
• 前馈：O(T \cdot d^2)

注意注意力那一项是 T^2——这是”一次算完所有位置”的代价。 第 3.4 节会看到，KV 缓存省的正好就是这一项。

3.2.5 回到案例：训练时模型看到的前文，全是真实的

现在可以指出案例的关键了。

训练时，位置 t 的输入是 x_1, \dots, x_t——这段文本的每一个字都是真实的。 模型在这个前提下被要求预测 x_{t+1}。它从来没见过”自己写错一个字之后的前文”。

用本模型实测：

真实文本上，模型对自己 argmax 的平均置信度 0.239　平均 NLL 2.203

0.239 意味着模型押的那个字，平均只有 24% 的概率——它自己也不太确定。 这很正常，字符级语言模型本来就难。

但注意另一组数：

自由生成时，模型对自己所选 token 的平均置信度 0.157　平均 NLL 2.135

在自由生成时，模型比在真实文本上更”自信”（NLL 2.135 < 2.203）。 而它写出来的东西是退化的重复文本。

这个反差是本章的第一个重要线索：模型的自信和它的正确性脱钩了。 原因就在于，生成时它看到的前文是它自己造的——一个训练时从未出现过的分布。 3.7 节会把这个现象定量算出来。

3.2.6 一个容易忽略的细节：哪些位置该算损失

式 式 3.2 里的求和 t = 1 到 T 看着很干净，但真实训练里有个坑： 一个 batch 里的样本长度不一样，短的要 padding。

本案例的语料是连续文本流，每次切 64 个字符，长度天然一致， 所以没这个问题。但真实场景（变长对话、变长文档）一定会有。

假设一个 batch 里有两条样本，长度分别是 5 和 3，padding 到 5：

样本 A：白 日 依 山 尽
样本 B：春 江 潮 <pad> <pad>

如果照式 式 3.2 算，模型会被要求预测 <pad>。 而 <pad> 不是真实内容，预测它毫无意义。更糟的是： 如果 padding 占的比例高，模型的梯度会被”预测 pad”这件事主导—— 它会学到一个完全没用的技能，而且这个技能会挤占真正有用的容量。

正确做法：用一个掩码把 padding 位置的损失置零。

\mathcal{L} = \frac{\sum_{t=1}^{T} m_t \cdot \big(-\log p_\theta(x_{t+1} \mid x_{1:t})\big)}{\sum_{t=1}^{T} m_t} \tag{3.4}

其中 m_t = 1 表示位置 t 是真实 token，m_t = 0 表示是 padding。

所以这个式子在说：分子只累加真实位置的损失，分母是真实位置的个数。

注意分母也变了——从 T 变成 \sum m_t。 如果忘了改分母，则 loss 会被 padding 的比例稀释， 你会以为模型训得比实际好。

一个真实的数字感受

假设一个 batch 里平均 60% 是 padding（在处理变长对话时很正常）。

• 正确实现：loss = 真实位置的平均损失，比如 3.0
• 分母忘了改：loss = 3.0 \times 0.4 = 1.2

1.2 看起来比 3.0 好得多，你会以为模型收敛得很快。 实际上什么都没变——只是分母算错了。

这个 bug 的危险性在于它让指标变好。所有”让指标变好”的 bug 都会被监控曲线掩盖过去，直到你拿模型去推理，发现它什么都没学会。

本书第 1 章遇到过同类问题：掩码写错让 loss 变好（4.634 vs 4.810）， 因为模型在偷看答案。“loss 异常地低”永远比”loss 异常地高”更值得警惕。

注记停下来想一想 ⑤

上面说 padding 位置的损失要置零。那么 padding 位置的输入呢？ 要不要也置零？

（答案：通常不用，但要小心。

输入端的 padding 一般用 <pad> 的嵌入，模型会把它当成一个普通的 token。 因为我们不计算 padding 位置的输出损失，所以它不影响梯度—— 但它会影响注意力。

具体说：位置 6 的查询会关注到位置 4、5 的 <pad>。 如果不管，模型会把注意力分给这些无意义的键。

这就是”注意力掩码”存在的原因——它和”损失掩码”是两件事： - 损失掩码：决定哪些位置的预测被计分（影响损失） - 注意力掩码：决定哪些位置可以被看到（影响前向）

很多实现在这里出错：只做了其中一个。第 1 章讲过的”掩码 bug” 就是注意力掩码写错了——症状是 loss 变好，因为模型偷看了不该看的位置。）

3.2.7 掩码怎么实现：−inf 还是大负数

式 式 3.3 是个数学定义。代码里怎么实现它，有一个真实的坑。

正确做法：填 -\infty。

att = att.masked_fill(mask == 0, float("-inf"))
att = F.softmax(att, dim=-1)

softmax 里 e^{-\infty} = 0，所以被屏蔽的位置权重严格为 0。

错误做法：填一个”足够小的数”。

att = att.masked_fill(mask == 0, -1e9)     # 看起来够小了
att = F.softmax(att, dim=-1)

e^{-10^9} 在 float32 下下溢成 0，所以结果通常是对的。

为什么”通常对”是危险的

-1e9 的写法有三个问题：

问题一：-1e9 加上别的东西会变。 如果 att 里有其他运算 （比如 ALiBi 的线性偏置、相对位置偏置），-1e9 + 5000 还是 -1e9， 不影响。但如果偏置本身量级到了 10^9（数值不稳定时会）， 就可能”抬起”这个值。

问题二：它破坏了 softmax 的数值稳定性技巧。 标准实现会先减去每行的最大值再取指数：

att = att - att.max(dim=-1, keepdim=True).values
att = att.exp()
att = att / att.sum(dim=-1, keepdim=True)

如果被屏蔽的位置是 -1e9，减完之后它变成 -10^9 - \text{max}， 仍然下溢成 0，没问题。但如果整行都被屏蔽（某个位置看不到任何东西）， max 就是 -1e9，减完变成全 0，softmax 输出均匀分布—— 一个本应是”全 0”的行变成了”全 1/4”。

问题三：-inf 会被显式地检查出来。 -inf 参与运算会产生 nan，而 nan 会立刻让训练崩掉—— 这是一个响亮的失败。 -1e9 产生的是静默的错误值。

又是”响亮失败 vs 静默降级”的权衡（见 3.6.3 节）。 本书的立场不变：宁可崩，不可悄悄错。

提示那”整行被屏蔽”怎么办

上面的问题二提了一个真实场景：某些位置看不到任何东西。

什么时候会发生？——padding。如果一个 batch 里某条样本特别短， 它后面的位置全是 padding。如果注意力掩码把这些位置全部屏蔽了， 那一行就没有任何有效键。

正确做法：不要把整行屏蔽光。 通常保留”位置可以看到自己” （对角线永远为 1），这样每行至少有一个有效键。

这也是 3.2.6 节那个”停下来想一想”的答案： 注意力掩码和损失掩码是两个独立的掩码， 而注意力掩码的设计要保证”每行至少有一个 1”。

3.3 推理那一面：一次只能算一个

3.3.1 自回归的死结

训练可以并行，推理不行。原因简单到有点扫兴：

要预测第 t+1 个字，需要第 t 个字的输出。而第 t 个字的输出， 需要第 t 个字本身。可第 t 个字还没生成。

这不是工程问题，是任务定义本身的性质。生成的定义就是”用已经生成的去生成下一个”， 所以它在时间上必然串行。

写成代码就是最朴素的样子：

out = prompt_ids
for _ in range(n_new):
    logits, _ = model(out)            # 把整段前缀重算一遍
    nxt = logits[:, -1, :].argmax(-1, keepdim=True)
    out = torch.cat([out, nxt], dim=1)

注意第 3 行：每一步都把整段前缀重新跑一遍前向。 第 1 步算 7 个位置，第 2 步算 8 个，第 3 步算 9 个……

这个写法有多浪费

生成第 t 个字的时候，前 t-1 个位置一个字都没变， 可它们的表示被从头算了一遍。

而且注意：在这个模型里，“位置 t 的表示”在生成过程中是永远不变的。 第 5 步算出来的位置 3 的表示，和第 4 步算出来的位置 3 的表示， 和最终第 50 步算出来的位置 3 的表示——完全一样。

为什么？因为因果掩码保证了位置 3 只看得到位置 1、2、3。 后面追加多少 token，都影响不到它。

所以这些计算不是”算得慢”，是纯粹的重复劳动。

3.3.2 两副面孔的计算量对比

图 3.2　 训练与推理的两副面孔。上：训练一次前向覆盖所有位置，注意力是 T×T 的矩阵，每个位置都有输出。下：推理一次只算一个位置，深色的格子是「这一步真正需要的」，浅色的是被朴素写法白白重算的。

把两边的计算量写出来，差距有多大就一目了然了。

生成 N 个字，提示词长度 P：

	每次前向处理的 token 数	前向次数	总 token 数
朴素推理	P, P+1, P+2, \dots, P+N-1	N	\sum_{i=0}^{N-1}(P+i)
训练式并行	P+N	1	P+N

本案例：P = 7，N = 50。

朴素推理的总 token 数 = \sum_{i=0}^{49}(7+i) = 50 \times 7 + \frac{49 \times 50}{2} = 350 + 1225 = 1575。

如果有一个”上帝视角”的写法能一次算完，只要 57 个 token。 1575 / 57 ≈ 27.6 倍。

实测（脚本里直接数）：

朴素（每步重算前缀）： 1,060 次 token 前向
带 KV 缓存        ：    46 次
减少 23.0 倍

（这里的 46 = 预填充 7 次 + 之后每步 1 次 × 39 步，比 57 少是因为 generate_cached 的第一次预测复用了预填充的输出。）

注记停下来想一想 ⑥

上面那个 \sum_{i=0}^{N-1}(P+i) 里，当 N 远大于 P 时， 它约等于 N^2/2。也就是说朴素推理的总计算量随生成长度平方增长， 而缓存版本是线性的 N。

如果生成 1000 个字（P 忽略不计），两者差多少倍？ 如果是 10000 个字呢？

（答案：500 倍和 5000 倍。注意这是个随长度增长的差距—— 短文本上感觉不到，长文本上是数量级的差别。）

3.3.3 推理的两个阶段：prefill 和 decode

在讲缓存之前，需要先分清推理的两个阶段。它们的瓶颈完全不同， 不分清就理解不了 3.4.4 节那个反直觉的实测结果。

阶段一：prefill（预填充）。 用户给了一段提示词，模型要把它”读进去”。 这一段是一次前向算完的——和训练时一模一样， 因为提示词的所有 token 都是已知的。

阶段二：decode（解码）。 逐字生成。每个新 token 都要一次前向， 而且必须等上一个算完。

两者的计算特征：

	prefill	decode
一次前向处理几个 token	P（提示词长度，可能几千）	1
矩阵形状	大而”方”	极瘦（1 \times d）
算术量	大	极小
瓶颈	算力（compute-bound）	显存带宽（memory-bound）
GPU 利用率	高（几十 %）	极低（个位数 %）

为什么 decode 是”显存带宽受限”

想象一次 decode 前向要做什么：把整个模型的所有权重 从显存读进计算单元，然后只做一次 1 \times d 的矩阵乘法。

用一个具体数字。本案例的模型有 425,984 个参数，float32 占 1.7 MB。 3090 的显存带宽约 936 GB/s，所以把 1.7 MB 读完要 1.7 \times 10^6 / 936 \times 10^9 \approx 1.8 微秒。

而算术量是 2 \times 425984 \approx 8.5 \times 10^5 次浮点运算， 在 35 TFLOPS 上是 0.024 微秒。

算得比读得快 75 倍。 也就是说，decode 阶段 GPU 有 98.7% 的时间 在等数据从显存里搬过来，计算单元基本闲着。

这就是为什么 KV 缓存省了 68 倍计算，墙钟只快 1.07 倍—— 省下来的都是本来就空闲的算力。

警告一个重要的推论：batch 是唯一的解药

如果 decode 的瓶颈是”读权重”，那么读一次权重应该服务尽量多的请求。

这正是真实推理服务的做法：把多个用户的请求拼成一个 batch， 一次前向算 B 个序列的同一个位置。

• B = 1：读 1.7 MB 权重，做 1 次乘法 → 浪费
• B = 64：读 1.7 MB 权重，做 64 次乘法 → 权重读取被摊薄了 64 倍

注意这和 KV 缓存是完全正交的两件事：

优化	省什么
KV 缓存	每次前向里的 token 数（算力）
Batching	每次前向里的 权重复用次数（带宽）

本章 3.4.4 节的实测里，batch 从 1 提到 64， 朴素写法的耗时几乎没变（0.082s → 0.083s），缓存写法的耗时也几乎没变 （0.083s → 0.084s）——因为在这个规模下，两者都还没撞到带宽墙。 模型太小，权重只有 1.7 MB，读它是微秒级的事。

把模型放大到 Llama-3-70B（140 GB），局面就完全不同了： 读一遍权重要 150 毫秒。那时 batch = 1 的 decode 是每秒 6 个 token， 而 batch = 64 能做到每秒 400 个 token——差 64 倍。

注记停下来想一想 ⑦

上面说 decode 阶段”算得比读得快”，所以 GPU 在等数据。

那么为什么 KV 缓存还是有用的？既然算力本来就空闲， 少算一点不是没意义吗？

（答案：因为 KV 缓存省的不只是算力，还包括显存读写。

朴素写法每一步要重新读入并写出 t 个位置的中间激活—— 这些激活比权重小得多，但随 t 增长。 到 t = 8000 时，注意力矩阵是 8000 \times 8000， 光是它就有 2.56 亿个元素，超过了模型权重本身的规模。

所以缓存省的是两块：算力（decode 阶段本来就富余）和 中间激活的读写（这才是真正卡住 decode 的东西）。

3.3.4 一个具体的形状清单

把本案例模型在 decode 阶段的每一步涉及到的张量形状列全。 这个清单在排查 bug 时非常有用——形状不对是最容易发现的错误。

张量	形状	说明
idx	(1, 1)	这一步的输入 token
x（残差流）	(1, 1, 128)	嵌入 + 位置
qkv 投影输出	(1, 1, 384)	3 \times 128，Q/K/V 拼在一起
q / k / v	(1, 1, 4, 32)	拆成 4 个头，每头 32 维
转置后的 q	(1, 4, 1, 32)	头维度提到前面
缓存的 k	(1, 4, t, 32)	t = 已缓存长度，逐步增长
打分 q @ kᵀ	(1, 4, 1, t)	一个查询对所有缓存的键
注意力权重	(1, 4, 1, t)	softmax 之后，和为 1
attn @ v	(1, 4, 1, 32)	加权求和
合并头	(1, 1, 128)	变回 d_model
logits	(1, 1, 1200)	词表上的分数

这张表里藏着一个关键信息

看”缓存的 k“那一行的形状：(1, 4, t, 32)。

只有时间维 t 在变，其他三维固定。 这就是为什么缓存是”追加”—— 新来的 k 形状是 (1, 4, 1, 32)，在 dim=-2 上拼上去， 时间维从 t 变成 t+1。

再看”打分”那一行：(1, 4, 1, t)。

注意最后那个 t——它随缓存增长。这就是 3.4.4 节说的 “中间激活随 t 增长”的具体体现：每步的注意力权重向量长度是 t， 不是常数。

而这个向量在朴素写法里长度是 t，在缓存写法里也是 t—— 缓存并没有让这个向量变小。它让别的东西变小了： 朴素写法要算 t \times t 的完整打分矩阵，缓存写法只算 1 \times t 的一行。

省掉的正是那 t 行。

提示用形状表排查 bug

3.9 节那六种 bug，有四种会立刻在形状上暴露：

bug	形状异常
dim 写错（-1 而不是 -2）	缓存变成 (1, 4, 1, t*32)，打分变成 (1,4,1,t*32)
no_cat	缓存长度永远是 1，不增长
pos_zero	形状正常，位置内容错（形状查不出来）
kv_mismatch	形状正常，对应关系错（形状查不出来）

后两种是最难查的——它们不违反任何形状约束。 这也是为什么 3.9.1 节说”要比 logits”： 形状检查能挡住一半的 bug，剩下的一半只能靠数值比对。

3.3.5 回到案例：确定性从哪来

现在回头看 3.1 节的循环。

推理是串行的，每一步的输入是上一步的输出。而贪心解码的每一步是确定性的： 同一个输入，argmax 永远给出同一个输出。

把这两件事合起来看：

x_{t+1} = \arg\max_v \; p_\theta(v \mid x_1, \dots, x_t) \tag{3.5}

所以这个式子在说：下一个字完全由前面的字决定，没有任何随机性。

这意味着整个生成过程是一个确定性动力系统：状态是”已经生成的文本”， 转移规则是 argmax。这类系统有一个众所周知的性质——状态空间有限时， 轨道最终必然进入周期循环。

模型的上下文只有 64 位，每一位有 1200 种可能，状态空间虽然巨大但有限。 一旦某段历史 x_{t-k}, \dots, x_t 重复出现过，后面的输出就会一模一样地重演。

案例里那个 8-gram 的重复，就是它的周期。

为什么”屏蔽坏 token”治不了本

屏蔽 <unk> 之后，输出变成「不見昔年年不見，不見春風吹。不見君不得， 不見君家。不見君不見，不見，不如何處？何時？不見君不見，不」。

<unk> 确实没了，循环也还在，只是周期从 8 变成了 4，循环单元从标点换成了「不見」。

这是一个重要的思维训练：屏蔽坏 token 是在改变状态空间的形状， 不是在改变动力系统的性质。argmax 仍然是确定性的， 有限状态空间仍然必然导致循环。你只是把循环的位置挪了挪。

真正要打破循环，必须让转移规则不再确定——这就是 3.5 节的采样。

提示一个可以自己做的实验

把提示词换成任意一段文字，用贪心生成 200 个字， 用 3.1.3 节的 degeneracy() 函数量一下 distinct-2 和最长重复 n-gram。

你会发现它几乎必然退化，而且退化程度和提示词关系不大。 这不是这个 42 万参数模型的问题——GPT-2、GPT-3 用贪心解码， 一样会退化。Holtzman 等人 2020 年那篇 The Curious Case of Neural Text Degeneration 讲的就是这件事，见 3.10 节。

3.3.6 既然必须串行，那怎么加速：四条路

自回归的串行性是任务定义带来的，无法消除。但“必须串行”不等于”每步都必须那么贵”。 工业界有四条路，它们针对的是不同的瓶颈。

第一条：KV 缓存。 减少每步的计算量。3.4 节详述。 注意它不减少步数——这是 3.4.4 节那个反直觉结果的根源。

第二条：批处理（batching）。 减少权重的重复读取。 3.3.3 节讲过：decode 阶段 GPU 在等数据，把 B 个请求拼一批， 读一次权重服务 B 个序列。步数不变，但每步的吞吐提高 B 倍。

第三条：投机解码（speculative decoding）。 减少步数。

这是最巧妙的一条。用一个小模型（draft model）先快速猜出接下来 5 个 token， 然后让大模型一次前向并行验证这 5 个猜测：

小模型猜：  春风 又 绿 江 南
大模型验：  春风 又 绿 江 南   ← 一次前向，5 个位置同时验证
             ✓    ✓   ✓   ✓   ✓

因为验证是一次前向算 5 个位置（就像训练那样并行）， 所以它只需要 1 步而不是 5 步。如果小模型猜对了 4 个， 就相当于用 1 步生成了 4 个 token——提速 4 倍。

猜错怎么办？从第一个错的位置截断，用大模型自己的结果。 这保证了输出分布和不用投机解码时完全一致——它是个无损加速。

投机解码正好踩中本章另一个 bug

回头看 3.9.4 节：no_mask 在单步解码时是恒真的。

投机解码恰恰是”一次前向算多个位置”——它把 T_{new} 从 1 变成 5。 所以：

no_mask 在 T_new=1 时 max|Δlogit| = 0.000e+00
         分块 4 个时 = 3.335

那个分块实验（3.335）测的就是投机解码的场景。

“验证”这一步如果掩码写错，被验证的 5 个位置会互相看到对方—— 第 2 个草稿 token 会看到第 3、4、5 个草稿。 验证结果全错，但程序不崩，只是接受率异常低（因为验证总是不通过）。

排查这类问题的人通常会怀疑小模型猜得不准， 而不会想到是大模型的验证路径错了。

第四条：换掉自回归。

前三條都在优化”怎么更快地串行”，第四条问的是能不能不串行。 这是一个活跃的研究方向，也是本书下篇的主题——扩散模型。 3.3.7 节展开。

四条路针对的瓶颈不同，要一起用

方法	减少什么	提速上限
KV 缓存	每步计算量	与上下文长度相关
批处理	权重的重复读取	与 batch size 成正比
投机解码	步数	与接受率成正比（通常 2–4 倍）
非自回归	串行性本身	理论上 1 步

注意”提速上限”那一列没有一个是无限的。 真实系统同时用前三条， 叠起来能到 10–50 倍——但每一步的串行依赖始终在那里。

3.3.7 如果把自回归去掉

自回归的根本代价是：N 个 token 需要 N 次串行前向。

能不能一次生成全部 N 个？ 这就是非自回归生成（non-autoregressive generation，NAR）的目标。

最直接的做法：模型一次输出 N 个位置的分布，每个位置独立采样。 问题是它必然产生不连贯的文本——因为位置之间没有任何协调机制。 “白日依山尽”可能被生成成”白日依山尽”（对）也可能”白月依水尽”（乱）， 因为每个位置独立决策。

后来的方法都是在”独立”和”串行”之间找中间点：

方法	独立程度	迭代次数
自回归	完全依赖	N
掩码预测（MaskGIT、Mask-Predict）	部分依赖	常数（如 10 次）
扩散语言模型	部分依赖	常数（如 50 次）
纯 NAR	完全独立	1

核心思路：先生成一个粗糙的完整序列，然后迭代细化。 每次细化都用当前完整序列作为条件——所以每一步都是”一次前向算所有位置”， 可以并行。

提示这就是下篇要讲的扩散模型

扩散模型的语言版本（diffusion language model）用的正是这个框架：

• 前向过程：把一句话逐步”加噪”——对文本来说就是把 token 逐步替换成 随机 token 或 [MASK]
• 反向过程：从纯噪声开始，逐步”去噪”成一句话
• 每次去噪都是一次前向，N 个位置同时算

迭代 50 次就得到结果——相比自回归的 N 次（N 通常是几百到几千）， 步数少一个数量级。

第 5 章会把这个框架的数学（前向过程的闭式解、ELBO、重参数化）完整推一遍。 这里只需要记住动机：自回归的串行性是它最大的成本， 而扩散提供了一条”用迭代换并行”的路。

它的代价是什么：迭代次数是个超参数，太少了质量差，太多了慢； 而且没有 KV 缓存可用——因为每一步的输入都完全不同， 不存在”只追加、不修改”的缓存。3.4.2 节那个恒等性证明的前提不成立。

注记停下来想一想 ⑧

上面说扩散语言模型”没有 KV 缓存可用”。

为什么第 3.4.2 节的证明在这里失效了？

（答案：证明的关键前提是”过去的 token 不会变”。

自回归生成时，前缀是定死的——生成第 t+1 个 token 时， 前 t 个永远不变。所以它们的 K/V 算一次就够了。

扩散的迭代细化每一步都在改所有的位置。 第 5 次迭代时位置 3 的 token 可能和第 4 次时不同， 所以位置 3 的 K/V 每一步都得重算。

“缓存”这个概念成立的前提是”被缓存的东西不再变”。 前提没了，缓存也就没了。）

3.4 KV 缓存：把两副面孔接起来

3.4.1 先算一遍：三个 token 的手工账

图 3.3　 KV 缓存随生成步数增长。蓝色「复用」表示该位置的 K/V 已经在缓存里，这一步不用重算；绿色「新」是这一步唯一要算的。因果掩码保证了蓝色格子永远不会变——这是缓存成立的前提。

在写任何代码之前，先拿最小的例子把账算清楚。这是本书的规矩。

设置：上下文长度 T = 3，模型维度 d = 4，单头（h = 1，d_{head} = 4）。 忽略前馈网络和归一化——它们不参与”哪些计算被重复”的讨论。

我们要生成第 4 个 token。假设前三步已经算过，现在看第四步。

朴素写法：把 x_1, x_2, x_3 重新喂一遍，拿到位置 3 的输出，再去预测 x_4。

用 W_Q, W_K, W_V \in \mathbb{R}^{4 \times 4} 表示投影矩阵，那么：

Q = X W_Q, \quad K = X W_K, \quad V = X W_V, \quad X \in \mathbb{R}^{3 \times 4} \tag{3.6}

Q, K, V 都是 3 \times 4。注意力输出是：

\text{Attn}(X) = \text{softmax}\left(\frac{QK^\top}{\sqrt{4}}\right) V \tag{3.7}

其中 QK^\top 是 3 \times 3，softmax 逐行做，最后乘 V 得到 3 \times 4。 我们要的是最后一行（位置 3 的输出），用它去算 x_4 的 logits。

现在把它写成”哪个位置用了哪些输入”的形式：

计算	用到的输入
Q 的第 3 行	x_3
K 的第 1、2、3 行	x_1, x_2, x_3
V 的第 1、2、3 行	x_1, x_2, x_3
QK^\top 的第 3 行	q_3 和 k_1, k_2, k_3
输出的第 3 行	第 3 行注意力权重和 v_1, v_2, v_3

关键观察：要算位置 3 的输出，需要的 K 和 V 是前三个位置全部的。 而如果我们上一步（生成 x_3 时）已经算过 k_1, k_2 和 v_1, v_2， 这一步只有 k_3, v_3 是新的。

所以缓存的做法是：

K_{\text{cache}} = [\,k_1; k_2; k_3\,], \qquad V_{\text{cache}} = [\,v_1; v_2; v_3\,] \tag{3.8}

每生成一个 token，只把它对应的那一行 k_t, v_t 追加到缓存末尾。

这一笔账省了多少

朴素写法：每步重新算整段的 Q, K, V。位置 3 这一步要算 3 行 K，3 行 V。

缓存写法：位置 3 这一步只算 1 行 K，1 行 V。

到第 t 步，朴素算 t 行，缓存算 1 行——省了 t 倍。 把 t 从 1 累加到 N，总省下的是 O(N^2) 级别的投影计算。

注意：Q 是不能缓存的。新位置必须算自己的查询向量。 但 Q 只有 1 行，K 和 V 是要被反复用到的——这就是为什么缓存的是 K 和 V。 名字也是这么来的。

注记停下来想一想 ⑨

上面说”Q 不能缓存”。请你从定义出发说明为什么。

（提示：q_t = x_t W_Q，它只依赖位置 t 自己的输入。 生成 x_{t+1} 时我们算的是位置 t+1 的 Q—— 这个位置以前根本不存在，所以没有”缓存”可言。 它是新的，不是重复的。

而 K、V 不一样：生成 x_{t+1} 时，注意力需要看到的是所有历史位置的 K 和 V。 这些历史位置的 K、V 在它们各自被生成时就算出来了，而且以后再也不会变。）

3.4.2 为什么它是恒等的：一个必须证明的前提

KV 缓存号称”不改变结果，只改变速度”。这是个很强的声明，得证明。

要证明的核心命题是：

位置 t 的 k_t 和 v_t，在生成第 t+1, t+2, \dots 个 token 的过程中， 永远等于最初算出来的那个值。

如果这个命题成立，缓存就是安全的——存下来的东西永远有效。 如果它不成立，缓存就是错的——存下来的是过期数据。

证明：

k_t 由什么决定？回顾注意力的计算：

k_t = \text{LN}(x_t) \, W_K \tag{3.9}

而 x_t 是位置 t 的残差流（residual stream）的值，它由第 t 个 token 的嵌入， 经过前面每一层的”注意力 + 前馈”累积而来：

x_t^{(\ell+1)} = x_t^{(\ell)} + \text{Attn}^{(\ell)}(x_1^{(\ell)}, \dots, x_t^{(\ell)}) + \text{FFN}^{(\ell)}(\cdot) \tag{3.10}

关键就在这个式子里：位置 t 在第 \ell+1 层的值，依赖于第 \ell 层里 位置 1 到 t 的值，不依赖任何 t' > t 的位置。

为什么？因为因果掩码。注意力在位置 t 只对 1 \dots t 加权求和， 那些 t' > t 的位置权重恒为 0。

归纳下去：x_t 只依赖 x_1, \dots, x_t 的初始嵌入， 也就是只依赖 token 1 \dots t。

而 token 1 \dots t 在生成过程中是不变的。 生成第 t+1 个 token 时， 前 t 个 token 早就定下来了，不可能再改。

所以 x_t 不变 \Rightarrow k_t, v_t 不变。命题得证。

这个证明为什么重要

它把”缓存”从一个工程技巧变成了一个数学推论。

工程技巧的意思是：这样做大概能行，如果结果有细微差别也正常。 数学推论的意思是：结果必须一模一样。差一点点就是 bug。

这个区别决定了我们怎么测试它。如果是技巧，测试方法是”看起来差不多就行了”； 如果是推论，测试方法是”逐位比对，看最大差异是不是 0”。

本章后面所有的缓存排错，都建立在这个测试方法上。

注记停下来想一想 ⑩

上面的证明用到了”因果掩码”这个前提。如果模型不用因果掩码 （比如 BERT 那样的双向注意力），缓存还成立吗？

（答案：不成立。双向注意力下，位置 t 的表示依赖于整句话的所有位置， 包括后面的。生成一个新 token 会改变所有旧位置的表示， 存下来的东西立刻过期。

这就是为什么只有自回归模型能缓存 K/V。 编码器模型（BERT）没有这个优化，因为它根本不做自回归生成。）

3.4.3 实测：0.0 和 1.9 × 10⁻⁶

证明完了，跑一遍验证。

测试一：预填充阶段。 把 7 个字的提示词一次喂进去（不带缓存），对比”一次喂进去但走缓存”：

预填充 max|Δlogit| = 0.0

完全相同的 0。不是 1e-7，不是 1e-9，是浮点数意义上的逐位相同。

测试二：增量阶段。 这是真正的考验——先缓存前 k 个位置，然后只喂第 k+1 个 token， 和不带缓存地重新跑一遍整段，对比第 k+1 个位置的 logits：

增量   max|Δlogit| = 1.907e-06

不是 0 了。 1.9 × 10⁻⁶。

这个非零值需要解释，它不代表缓存错了。

警告1.9e-6 是什么

浮点数加法不满足结合律。

(a + b) + c 和 a + (b + c) 在实数上相等，在 float32 上可能差最后一位。 缓存路径和重算路径，对同一批数做加法的顺序不一样： - 重算路径：一次矩阵乘法，内部按某种分块顺序累加 - 缓存路径：历史部分是一次乘法，新 token 部分是另一次，然后把结果相加

顺序不同，最后一位就可能不同。1.9 × 10⁻⁶ 相对于 logits 的量级（约 10）是 相对误差 2 × 10⁻⁷，正好是 float32 的机器精度（约 1.2 × 10⁻⁷）量级。

判断标准：这个差异应该随”模型规模”和”序列长度”缓慢增长， 但不应该随”生成步数”累积爆炸。如果它随步数指数增长，那是真的错了。

提示一个实用的判据

实践中怎么区分”浮点误差”和”真的写错了”？看量级。

观察到的 max|Δlogit|	判断
< 10^{-4}	浮点误差，实现是对的
10^{-3} \sim 10^{-1}	可疑，要查
> 1	一定是 bug

本章实测的四种真 bug，最大差异分别是 12.06、10.85、17.53、2.10—— 最小的那个 2.10 也比浮点误差大六个数量级。

这个判据不是拍脑袋定的：logits 本身的量级在 10 上下， float32 有 24 位有效数字，10 \times 2^{-24} \approx 6 \times 10^{-7}。 累积几步之后 10^{-5} 是合理上限。

3.4.4 省了 28 倍计算，为什么一点没变快

图 3.4　 KV 缓存省下的计算量（上条）与实测墙钟加速比（下条）。同一组里两条相差几十倍——省掉的算术没有变成省掉的时间，因为 decode 阶段的瓶颈是 kernel 启动和显存带宽，不是算力。

现在到了本章最反直觉的部分。

既然缓存省掉了 23 到 28 倍的计算量，那它应该快 23 到 28 倍吧？

实测结果：

cuda:1 batch=  1：朴素 0.082s　缓存 0.083s　实测  0.99×　计算量 28.1×
cuda:1 batch= 16：朴素 0.105s　缓存 0.101s　实测  1.04×　计算量 28.1×
cuda:1 batch= 64：朴素 0.083s　缓存 0.084s　实测  0.99×　计算量 28.1×
   cpu batch=  1：朴素 0.049s　缓存 0.034s　实测  1.44×　计算量 28.1×
   cpu batch= 16：朴素 0.137s　缓存 0.054s　实测  2.53×　计算量 28.1×
   cpu batch= 64：朴素 0.385s　缓存 0.077s　实测  4.97×　计算量 28.1×

GPU 上，缓存一点没快。 计算量明明少了 28 倍。

你可能会说：是不是模型太小了？42 万参数确实小。那就换个大的：

── 大模型·8层768（58,387,968 参数）
   cuda:1 生成  32：朴素   0.179s　缓存  0.187s　实测  0.95×　计算量  19.3×
   cuda:1 生成  64：朴素   0.367s　缓存  0.358s　实测  1.03×　计算量  35.6×
   cuda:1 生成 128：朴素   0.766s　缓存  0.713s　实测  1.07×　计算量  67.8×
      cpu 生成  32：朴素   0.773s　缓存  0.514s　实测  1.51×　计算量  19.3×
      cpu 生成  64：朴素   1.774s　缓存  1.015s　实测  1.75×　计算量  35.6×

5839 万参数、8 层、计算量省了 68 倍，GPU 上实测 1.07 倍。

为什么

因为墙钟时间不由计算量决定，由”有多少步必须串行”决定。

看两条路径的结构：

	前向次数	每次前向的 token 数
朴素	N	P, P+1, \dots（递增）
缓存	N（预填充 1 次 + 每步 1 次）	P，然后全是 1

前向次数是一样的，都是 N 次。差别只在每次前向算多少 token。

而一次前向在 GPU 上要花多少时间，绝大部分不是花在算术上， 而是花在：Python 解释器的开销、kernel 启动、显存读写、同步。

算一下单次前向的量级。8 层、d = 768、d_{ff} = 3072， 一层的前馈网络约 2 \times d \times d_{ff} = 2 \times 768 \times 3072 \approx 4.7 \times 10^6 次乘加，注意力约 2 \times T \times d^2。取 T = 128，一层合计约 10^8 量级， 8 层约 10^9。本机 3090 的 fp32 算力约 35 TFLOPS， 所以纯算术时间约 30 微秒。

而实测单步耗时约 5.5 毫秒。99% 的时间不在算术上。

这个结论的正确用法

不是说”KV 缓存没用”。而是说它省的是算力，不是延迟。

这两件事在不同的场景下价值完全不同：

场景	瓶颈	KV 缓存的价值
单条请求、短上下文、GPU	kernel 启动 + Python 开销	几乎为零
单条请求、CPU	算力	1.4–5.0 倍
大批量服务	显存带宽 + 算力	显著
长上下文（T 很大）	注意力的 O(T^2)	决定性

最后一行是关键。本章的上下文只有 1024，注意力的 T^2 项还是小头。 真实模型动辄 128K 上下文——那时朴素写法每一步都要算 T^2 的注意力矩阵， T = 128000 时这个矩阵有 1.6 \times 10^{10} 个元素， 根本跑不起来。KV 缓存不是”优化”，是”能不能跑”的区别。

注记停下来想一想 ⑪

既然 GPU 上瓶颈是”前向次数”而不是”每次算多少”， 那么一个自然的推论是：把多次前向合并成一次应该能提速。

这个想法是对的吗？想想为什么现在的推理框架要做 “continuous batching”（把不同请求的 token 拼成一个 batch 一起算）。

（答案：对。continuous batching 干的就是这件事—— 它不减少总计算量，但把多个请求的”第 k 步”合并成一次前向， 于是前向次数按并发数摊薄。这和 KV 缓存省计算量是两个正交的优化方向， 实践中要一起用。）

3.4.5 缓存的形状与成本

缓存不是免费的。把它的开销算清楚。

显存占用：

\text{bytes} = 2 \times L \times h \times T \times d_{head} \times \text{sizeof(dtype)} \tag{3.11}

所以这个式子在说：2（K 和 V 各一份）× 层数 × 头数 × 上下文长度 × 每头维度 × 数据类型字节数。

注意 h \times d_{head} = d_{model}，所以可以简化成：

\text{bytes} = 2 \times L \times T \times d_{model} \times \text{sizeof(dtype)} \tag{3.12}

本案例的模型：L = 2，T = 64，d_{model} = 128，float32（4 字节）。

2 \times 2 \times 64 \times 128 \times 4 = 131{,}072 \text{ 字节} \approx 128 \text{ KB}

比模型本身（425,984 个参数 × 4 = 1.7 MB）还小。 所以这里完全不用考虑显存。

换成真实规模算一下，感受会完全不同。Llama-3-70B：L = 80， d_{model} = 8192，T = 8192，fp16（2 字节）：

2 \times 80 \times 8192 \times 8192 \times 2 \approx 21.5 \text{ GB}

单个请求、单个序列，光 KV 缓存就要 21.5 GB。 模型权重本身是 140 GB。缓存占到了权重的 15%——这就是为什么真实推理服务 要在”并发多少条请求”和”每条多长上下文”之间反复权衡。

注记停下来想一想 ⑫

式 式 3.12 里，T 是线性出现的。但注意 T 在这台机器上 等于”已经生成了多少 token”——它会随着生成过程不断增长。

这意味着什么？

（答案：KV 缓存的显存占用是动态增长的。一条请求开始时只占几百 KB， 生成 8000 个 token 后就占满了。这也是为什么推理框架要做 “显存预分配 + 分页管理”（vLLM 的 PagedAttention 就是干这个的）—— 否则每生成一个 token 都要重新分配一块更大的显存，还要把旧数据拷过去。）

3.4.6 回到案例：缓存解决了问题的一半

把这一节接回 3.1 节的案例。

缓存解决了什么？它让”每一步只算新 token”成为可能，从而 让生成的总计算量从 O(N^2) 降到 O(N)。

但注意——它没有解决循环问题。

none        首个 logit 差在第 None 步　max|Δ|=0.0000   　文本分歧  0 处

缓存版本生成的文本，和不带缓存一字不差。它还是「。，。，。，」。

这正是”数学推论”的直接后果：如果缓存是恒等变换， 它必然产生完全相同的输出，包括完全相同的错误。

要治循环，得动解码策略，不是动缓存。 这是 3.5 节的事。

3.4.7 缓存的工程变体：从 MQA 到分页

3.4.5 节算过：Llama-3-70B 单序列、8K 上下文，KV 缓存要 21.5 GB。 这一节看工业界怎么把它压下来。

变体一：MQA —— 所有头共享一份 K/V

标准多头注意力里，每个头有自己的 K、V。h = 64 个头就有 64 份缓存。

MQA（multi-query attention）的做法：Q 还是每个头一份，K/V 全局只有一份。 缓存大小直接降到 1/h——70B 模型上从 21.5 GB 降到 0.34 GB。

代价是质量下降。因为不同头本来要关注不同的东西 （第 2 章讲过：多头买的是”让不同头学不同关系”的自由度）， 共享 K/V 之后这个自由度就没了。

变体二：GQA —— 分组共享

GQA（grouped-query attention）是折中：把 h 个头分成 g 组， 组内共享 K/V，组间不共享。

• g = h：退化成标准 MHA
• g = 1：退化成 MQA
• g = 8：Llama-2/3 的选择，缓存降到 1/8，质量损失很小

这是个漂亮的工程折中：没有理论上的最优解， 但实证发现”8 组”基本无损。

注意它和本章 bug 的关系

MQA/GQA 改的是缓存的形状：从 (B, h, T, d_head) 变成 (B, g, T, d_head)。

它不改变缓存的语义——仍然是”只追加、不修改”， 仍然满足 3.4.2 节那个恒等性证明的前提。

所以 3.9.1 节那套”比 logits”的测试方法，换成 MQA 之后完全适用。 你不需要为 MQA 写一套新的测试——只要把形状的期望值改掉。

这是好抽象的标志：一个改变不会波及测试方法论。

变体三：分页 —— 借操作系统的思路

前面两个变体压缩的是每个 token 的缓存大小。分页解决的是另一个问题： 缓存的总量管理。

问题的根源在 3.4.5 节的”停下来想一想 ⑦”：缓存的显存占用是动态增长的， 而且生成之前你不知道会长到多少。

传统做法是预分配一个最大长度的块（比如 8192）。后果：

• 一条只生成 100 个 token 的请求，占了 8192 的空间
• 实测中浪费率高达 60%–80%

vLLM 的 PagedAttention 借用操作系统虚拟内存的思路： 把缓存切成固定大小的页（比如 16 个 token 一页）， 用一张页表把逻辑位置映射到物理页。

好处有三个： 1. 按需分配：用到第几页才分配第几页 2. 无外部碎片：所有页大小相同，随便放 3. 前缀共享：同一个提示词的多个输出共享同一份前缀页

第 3 点是意外收获，也是最大的收益之一。 很多应用（few-shot 提示、多轮对话）有大量重复前缀， 共享之后显存占用能降一个数量级。

前缀共享正好踩中本章一个 bug

回到 3.9.4 节那个结论：no_mask 在单步解码时是恒真的， 因为查询位置恰好是最大的键位置。

前缀共享破坏了这个前提。

设想两个请求共享了长度 1000 的前缀，然后各自继续生成。 请求 A 生成了第 1001 个 token，请求 B 生成了第 1002 个。 如果把 B 的缓存和 A 的拼接……查询位置和最大键位置不再相等， 掩码不再是恒真的。

所以那个”在我们这里不会出问题”的 bug， 会在引入前缀共享的那一天突然爆发。

这正是 3.9.4 节那句话的意思：理论上不会发生的事情，实践中总会发生一次。

3.4.8 缓存里到底存了什么：一次完整的手工推演

3.4.1 节算的是”省了多少”，这一节把”每步到底算什么”完整走一遍。

设置：上下文长度 4，d_{model} = 4，单头。提示词是 3 个 token： x1 x2 x3，现在要生成 x4。

第一步：预填充。

输入的 X_1 \in \mathbb{R}^{3 \times 4}（三个位置，每个 4 维）。 算出的：

Q_1 = X_1 W_Q \in \mathbb{R}^{3 \times 4}, \quad K_1 = X_1 W_K \in \mathbb{R}^{3 \times 4}, \quad V_1 = X_1 W_V \in \mathbb{R}^{3 \times 4} \tag{3.13}

打分矩阵 S_1 = Q_1 K_1^\top / \sqrt{4} \in \mathbb{R}^{3 \times 3}。 加因果掩码（下三角）后 softmax，得到 A_1 \in \mathbb{R}^{3 \times 3}。 输出 A_1 V_1 \in \mathbb{R}^{3 \times 4}。

我们只要最后一行 A_1[3,:] V_1——它是位置 3 的输出，用来预测 x4。

把 K、V 存起来：K_{\text{cache}} = K_1，V_{\text{cache}} = V_1， 形状都是 3 \times 4。

第二步：生成 x4。

x4 被喂进来，算出它自己的：

q_4 = x_4 W_Q \in \mathbb{R}^{1 \times 4}, \quad k_4 = x_4 W_K, \quad v_4 = x_4 W_V \tag{3.14}

关键：q_4 是新的（位置 4 以前不存在），k_4, v_4 也是新的。 但注意力需要的是”位置 4 对所有历史位置的关注”—— 所以需要 k_1, k_2, k_3，而它们已经在缓存里了。

拼接：

K_{\text{cache}}' = [k_1; k_2; k_3; k_4] \in \mathbb{R}^{4 \times 4}, \quad V_{\text{cache}}' = [v_1; v_2; v_3; v_4] \tag{3.15}

打分：S = q_4 K_{\text{cache}}'^{\top} / \sqrt{4} \in \mathbb{R}^{1 \times 4}。

掩码：这里出现了一个 3.4.1 节没提的细节。

朴素路径的掩码是 4 \times 4 的下三角。缓存路径的掩码是 1 \times 4—— 只有一行，因为只有一个查询。

这一行是什么？ 查询位置是 4，键位置是 1、2、3、4。 因果规则”键位置 \le 查询位置”对四个键全部成立。 所以这一行是 [1, 1, 1, 1]——全是 1，什么都不屏蔽。

这就是 3.9.4 节那个结论的来源： 在单步解码时，因果掩码恒真。 不是”掩码不重要”，而是”这一步恰好没有未来需要屏蔽”。

输出：A V_{\text{cache}}' \in \mathbb{R}^{1 \times 4}，用它预测 x5。

把两次的形状并排看

	预填充	单步解码
查询数	3	1
键数	3	4
打分矩阵	3 \times 3	1 \times 4
掩码	下三角	全 1
新算的 K/V	3 行	1 行
复用的 K/V	0 行	3 行

“复用了 3 行”就是缓存的全部内容。 朴素写法会把这 3 行重算一遍。

到第 t 步，复用的行数是 t-1。累加起来，整个生成过程省下的 是 O(N^2) 行的投影计算——这就是 3.3.2 节那个 27.6 倍的来源。

注记停下来想一想 ⑬

上面说单步解码的掩码是”全 1”。那么如果模型没有因果掩码， 单步解码时还需要掩码吗？

（答案：在缓存路径下不需要，但那是因为别的原因。

如果不加因果掩码，预填充阶段就会错——位置 1 会看到位置 2、3。 所以预填充必须有掩码。

而单步解码时，缓存里的所有 token 都确实是过去（它们已经被生成了）， 所以”看到全部”是正确的行为——恰好等于全 1 掩码。

换句话说：单步解码的正确行为，恰好就是”不掩码”。 这个巧合是 3.9.4 节那个 bug 能藏住的根本原因。）

3.4.9 怎么给缓存写测试

3.4.2 节证明了缓存是恒等变换，3.4.3 节实测了它。这一节把它变成可复用的测试。

测试一：恒等性（最重要的一个）

def test_cache_is_identity(model, ids):
    """缓存必须逐位等于重算。"""
    ref, _ = forward_step(model, ids, cache=None)         # 整段一次算
    cache = KVCache(len(model.blocks))
    got, _ = forward_step(model, ids, cache=cache)        # 走缓存
    assert torch.equal(ref, got), f"预填充不一致：{(ref-got).abs().max()}"

注意用了 torch.equal 而不是 allclose。

为什么？因为预填充阶段两次计算走的是完全相同的代码路径 （都是 T 个位置一次算完），所以结果应该逐位相同。 用 allclose 会掩盖掉”其实走了不同路径”这个事实。

本章实测：max|Δlogit| = 0.0。 不是 1e-7，是 0。

测试二：增量一致性

def test_incremental(model, ids, k):
    """第 k 个位置：整段重算 vs 只算新 token。"""
    seq = ids[:, :k+1]
    full, _ = forward_step(model, seq, cache=None)
    cache = KVCache(len(model.blocks))
    forward_step(model, ids[:, :k], cache=cache)
    step, _ = forward_step(model, ids[:, k:k+1], cache=cache)
    d = (full[:, k] - step[:, 0]).abs().max().item()
    assert d < 1e-4, f"增量路径不一致：{d}"

这里用容差而不是 torch.equal，因为增量路径和整段路径的 加法顺序确实不同（见 3.4.3 节），浮点误差是预期的。

本章实测：1.907e-06，远小于 1e-4 的阈值。

两个测试的严格程度不同，这是有意的

测试	判据	为什么
预填充	torch.equal（严格）	两条路径应该完全一样
增量	< 1e-4（容差）	两条路径的加法顺序不同

如果两个都用严格相等：增量测试会失败（浮点误差）， 你会以为实现有 bug，然后去改一个正确的实现。

如果两个都用容差：预填充测试会漏掉”走了不同代码路径”的问题。

为每个测试单独定判据，而不是用一个统一的阈值—— 这需要你想清楚”这两条路径在数学上应该有多接近”。 测试的严格程度是设计出来的，不是抄来的。

测试三：故意注入错误，验证测试能抓住

@pytest.mark.parametrize("bug", ["pos_zero", "no_cat", "no_mask", "kv_mismatch"])
def test_bugs_are_caught(model, ids, bug):
    """把每种 bug 注进去，上面的测试必须失败。"""
    with pytest.raises(AssertionError):
        run_identity_test(model, ids, bug=bug)

这是最容易被跳过、却最重要的一步。

一个测试如果不管代码怎么写都通过，那它没有在测任何东西。 本章的五种 bug 实现（kv_cache.BUGS）存在的意义就是这个—— 它们是测试的测试。

实测它们全部被抓到（3.9.1 节的表：最大差异 2.10 到 17.53， 全部远超 1e-4 的阈值）。

提示测试覆盖率不等于测试有效性

这里有一个反直觉的事实：

本章的缓存实现只有 20 行代码，测试代码比实现代码长。

这在”数学上应该恒等”的场景里是正常比例。 因为这类代码的特点是：写对了不难，写错了也不报错。

所以投入产出比完全不一样： - 普通业务代码：测试是”防止未来改坏” - 恒等性优化代码：测试是”证明现在就是对的”

第二类的测试不是可选项，是交付物的一部分。

3.4.10 批处理时的缓存：一个绕不开的麻烦

前面所有讨论都假设 batch = 1（一条请求）。真实服务要同时处理很多条。

最初的直觉：把缓存从 (1, h, T, d) 变成 (B, h, T, d) 不就行了？

不行。因为每条请求的长度不一样。

设想 batch 里有两条请求：

请求 A：已生成 100 个 token
请求 B：已生成  10 个 token

现在要同时算它们的”下一步”。A 的查询位置是 100，B 的是 10。 它们的缓存长度不同，没法直接拼成一个张量。

三种解法，各有代价：

做法	怎么处理	代价
补零对齐	按最长的补到 100	B 的 90 个位置是 padding，白算
分组（bucketing）	长度相近的放一组	组内仍要补齐；组间负载不均
分页（PagedAttention）	物理不连续，用页表索引	需要自定义 kernel

第三种是现在的答案。 这是 3.4.7 节讲的 PagedAttention 的第二个动机—— 第一个动机是”避免预分配浪费”，第二个是”让变长请求能真正并行”。

数学上还有一个不显然的问题

补零对齐不只是”白算”——它会污染结果。

假设 B 的缓存被补到 100 个位置，后 90 个是零向量。 那么 B 的查询会算 \text{softmax}(q K^\top)，其中 K 包含 90 个零行。

零行的打分是 q \cdot 0 = 0。而 0 是一个合法的 logit！ 经过 softmax，这 90 个假位置会分到实实在在的概率质量。

结果：B 的注意力被 90 个不存在的 token 稀释了。

正确做法是给这些位置加掩码，让它们的权重严格为 0。 这就回到了 3.2.7 节那个”掩码怎么实现”的问题—— 而这次是被 padding 逼出来的，不是被因果性逼出来的。

同一段代码要处理两种不同来源的掩码： 1. 因果掩码：防止看到未来 2. padding 掩码：防止看到补零

很多实现只做了第一个。 在 batch = 1 时看不出问题 （没有 padding），一上 batching 就出问题。

注记停下来想一想 ⑭

上面说”零向量的打分是 0，经过 softmax 会分到概率质量”。

为什么这不影响 A？ A 的缓存里没有补零。

（答案：因为在”补零对齐”的做法里，A 不需要补—— 它是 batch 里最长的那条。B 才需要补。

所以症状是”只有部分请求出错”： 每条请求的输出质量取决于”它在 batch 里的相对长度”。 同一个提示词，单独跑是对的，和一条长请求放一起跑就变差了——

这种”依赖 batch 组成”的行为是最难排查的 bug 之一， 因为它不可复现：换个请求顺序、换个并发量，结果就变了。

这又是”静默降级”（见 3.6.3 节）：不报错，只是某些请求悄悄变差。

3.4.11 一句话总结缓存和批处理的分工

	KV 缓存	批处理
减少的量	每步的 token 数	每步的前向次数（摊销）
瓶颈类型	算力	显存带宽
在 CPU 上	明显有效（1.4–5.0 倍）	有效
在 GPU 上（小模型）	几乎无效	有效
在 GPU 上（长上下文）	决定性	有效
引入的新问题	显存管理、变长对齐	变长对齐、调度

两者必须一起用，而且两者都引入了”变长对齐”这个问题—— 这个问题的解法是同一个：分页。

3.4.12 一个自然的问题：训练时能用缓存吗

不能，而且不是因为”不需要”，是因为”用不上”。

3.2.3 节讲过：训练时一次前向处理整段 T 个位置， 每个位置只算一次。缓存的作用是”避免重复计算同一个位置”—— 训练时根本没有重复。

但有一个例外值得一提：多轮对话的训练。

假设一个训练样本是 10 轮对话，每轮平均 200 个 token，共 2000 个。 朴素做法是：把整个 2000 token 的序列喂进去，一次前向算完。

这没问题。但如果数据格式是”10 个独立的训练样本，每个都是 前缀 + 一轮回复”（这是常见的数据构造方式），那前缀被重复计算了 10 次。

这时有两种做法：

做法	计算量	梯度正确性
每个样本独立算	10 次前向，前缀算 10 遍	正确
拼成一个长序列	1 次前向	会出错

为什么”拼成一个长序列”会出错

因为不同轮次的回复之间不应该互相看见。

第 3 轮的回复不应该影响第 5 轮的前缀表示—— 但在一个长序列里，因果掩码会让第 5 轮看到第 3 轮的回复。

这不一定错——如果数据本来就是连续对话，那是对的。 但如果数据是”同一个前缀 + 10 个不同回复”（比如 RLHF 里的 同一 prompt 的多个采样），拼起来就是错的： 它把 10 个平行的样本变成了 1 个串行的样本。

正确做法：用块对角掩码（block-diagonal mask）—— 让每一轮只能看到”前缀 + 自己”，看不到其他轮次。

M = \begin{pmatrix} M_1 & 0 & 0 \\ 0 & M_2 & 0 \\ 0 & 0 & M_3 \end{pmatrix} \tag{3.16}

所以这个式子在说：整张掩码分成几个对角块， 块内是正常的因果下三角，块间全是 0。

训练时缓存用不上，但”怎么组织多个样本”这个问题绕不开， 而它的解法和掩码有关——和本章 3.2.7 节、3.9.4 节讲的是同一件事。

注记停下来想一想 ⑮

上面的块对角掩码，形状上是什么样子？

如果只有 2 个块、每块 2 个位置，把 4×4 的掩码画出来。

（答案：

M = \begin{pmatrix} 1 & 0 & 0 & 0 \\ 1 & 1 & 0 & 0 \\ 0 & 0 & 1 & 0 \\ 0 & 0 & 1 & 1 \end{pmatrix}

注意它不再是下三角矩阵。 位置 3 看不到位置 1、2—— 这是块间为 0 的效果。

这意味着”因果掩码 = 下三角”这个简化说法只在单样本时成立。 真实训练代码里的掩码是构造出来的，不是简单取 tril。

如果你在别人的代码里看到 torch.tril，那说明那个实现 只支持单样本训练。 支持多样本打包的实现一定有更复杂的掩码构造逻辑。）

3.5 采样的数学：让输出不再确定

3.5.1 从 argmax 到”按概率抽”

3.3.5 节证明了：贪心解码是一个确定性动力系统，有限状态空间下必然循环。

最直接的解法是把确定性去掉。不是取概率最大的那个字，而是按概率抽签：

x_{t+1} \sim \text{Categorical}\big(p_\theta(\cdot \mid x_1, \dots, x_t)\big) \tag{3.17}

所以这个式子在说：每个字被选中的概率，就等于模型给它的概率。 概率 0.3 的字，抽 100 次大约中 30 次。

循环还会发生吗？理论上还会（任何有限状态马尔可夫链都有常返态）， 但概率极低：要重现循环，需要连续几十步都恰好抽中那一个特定的字。 如果那个字的概率是 0.5，连续 8 次都抽中的概率是 0.5^8 = 0.4\%。 在贪心解码下，这个概率是 100%。

从 100% 到 0.4%，但没有到 0

实验数据印证了这一点：

解码	distinct-2	最长重复
贪心	0.102	8-gram
τ = 0.8	1.000	无

但注意”无”是在这 50 个 token 里没有出现。多生成一些， 重复片段一定会出现——只是不会再出现”连续重复 8 次”这种病态情况。

采样不是消灭了重复，是把病态重复变成了正常重复。 真实文本里也有重复的字词，只是不会像复读机那样卡死。

但是直接用模型原始概率采样不行。原因在实验数据里：

τ=1.0  熵  2.822 nats　困惑度   16.8　最大概率 0.4529

模型原始输出的最大概率只有 0.45。剩下的 0.55 摊在 1199 个字上。 这意味着：每次都有一半以上的概率抽到一个”模型并不想写”的字。

抽 50 次，几乎必然抽到几个完全不着调的字，而且一旦抽到， 后面就顺势跑偏了。这就是”纯采样”输出会语无伦次的原因。

所以我们需要的不是”完全随机”，而是“在一个受控的候选集里随机”。 接下来的三种技术，都是在定义这个候选集。

3.5.2 温度：把分布重新标定一次

图 3.5　 温度对真实模型分布的影响。青色是熵（左轴），橙色是最大概率（右轴）。τ<1 让分布变尖，熵快速塌陷；τ>1 让分布变平，但熵会饱和在 ln(1200)=7.09。

温度做的事，是在 softmax 之前把 logits 除以一个数：

p_v^{(\tau)} = \frac{\exp(z_v / \tau)}{\sum_{u} \exp(z_u / \tau)} \tag{3.18}

所以这个式子在说：\tau < 1 时，所有 logits 被放大， 原本的差距也被放大，概率分布变得更尖；\tau > 1 时被压平，分布更均匀。

两个极端：

• \tau \to 0^+：除以一个趋于 0 的正数，logits 差距趋于无穷。 最大的那个 logit 独占全部概率。温度采样退化成 argmax。
• \tau \to \infty：所有 logits 都趋近于 0，分布趋近于均匀分布。 温度采样退化成词表上的均匀随机。

用实验里那一组数走一遍

模型在某一步给出的 logits，我们取它最大的三个分量做例子： z = [4.0, 3.0, 2.0]（其余 1197 个分量更低）。

τ = 1.0：e^4 = 54.60，e^3 = 20.09，e^2 = 7.39，和 = 82.08。 概率 [0.665, 0.245, 0.090]。最大项 0.665。

τ = 0.5：logits 变成 [8.0, 6.0, 4.0]。 e^8 = 2981，e^6 = 403.4，e^4 = 54.6，和 = 3439。 概率 [0.867, 0.117, 0.0159]。最大项从 0.665 涨到 0.867。

τ = 2.0：logits 变成 [2.0, 1.5, 1.0]。 e^2 = 7.39，e^{1.5} = 4.48，e^1 = 2.72，和 = 14.59。 概率 [0.506, 0.307, 0.186]。最大项从 0.665 降到 0.506。

同一个模型、同一组 logits、同一批候选字，只是换了个 τ， “最可能写哪个字”的把握就从 0.51 到 0.87。

实测：温度对真实分布的影响。

在提示词”春江潮水连海平”之后那一步，模型给出一个 1200 维的分布。 它的熵（不确定性的度量）随 τ 变化如下：

τ	熵（nat）	困惑度	最大概率
0.1	0.003	1.0	0.9996
0.3	0.245	1.3	0.9338
0.5	0.519	1.7	0.8236
0.8	1.489	4.4	0.6323
1.0	2.822	16.8	0.4529
1.5	5.442	231.0	0.1393
2.0	6.314	552.0	0.0514
3.0	6.790	889.2	0.0153

这张表里有两个要看的现象

第一，熵不是线性增长的。 τ 从 0.1 到 1.0（10 倍），熵从 0.003 到 2.822（940 倍）； 但从 1.0 到 3.0（3 倍），熵只从 2.822 到 6.790（2.4 倍）。

为什么？ 因为 \tau \to \infty 时熵的上限是 \log V = \log 1200 = 7.09 nat。 表里 τ = 3.0 时熵是 6.790，已经逼近上限了——继续加温度也只是趋近均匀分布， 不会无限增长。

这个上限的存在意味着：“温度调高”的效果会饱和。 τ = 3 和 τ = 5 的输出，随机程度差不多。

第二，τ < 1 时熵塌陷得极快。 τ = 0.3 时熵只有 0.245 nat， 困惑度 1.3——相当于模型在 1.3 个字里选。这已经接近贪心了。

这解释了实践中的经验：温度的有效区间很窄。 大模型常用的值是 0.7 到 1.0；低于 0.5 基本等于贪心，高于 1.5 就开始胡说。

注记停下来想一想 ⑯

表里 τ = 0.1 时，最大概率是 0.9996，熵是 0.003 nat。

如果 \tau = 0 会怎样？\tau 能取 0 吗？

（提示：式 式 3.18 里 \tau 在分母上。\tau = 0 是除零， 代码里会得到 inf 或 nan。

所以工程实现里没有”温度 = 0”这个选项。所谓的”温度 0”是一个约定， 它的实际含义是”跳过采样，直接 argmax”。很多 API 文档会写 temperature: 0 表示确定性输出——那是接口层面的语法糖，不是数学。）

3.5.3 top-k：只在前 k 个里抽

温度的问题是它改变了整个分布的形状，包括那些概率极低的尾部。 τ = 0.8 时，第 500 个候选字仍然有非零概率——而它可能是个完全不相干的字。

top-k 的思路更直接：只保留概率最大的 k 个字，其余置零，重新归一化。

p'_v = \frac{p_v \cdot \mathbb{1}[v \in S_k]}{\sum_{u \in S_k} p_u} \tag{3.19}

其中 S_k 是概率最大的 k 个字的下标集合。

所以这个式子在说：砍掉尾巴，把砍剩下的重新归一化。 分母 \sum_{u \in S_k} p_u 是”被保留的那部分概率之和”—— 除以它是为了让新分布仍然加起来等于 1。

一个真实的例子

从实测数据看，τ = 1.0 时这个分布的最大概率是 0.4529。 如果 top-k = 10，那么保留的是概率最高的 10 个字。

问题来了：这 10 个字加起来占多少概率？

我们的实测（τ = 1.0）显示，累积概率达到 0.5 需要 2 个字， 达到 0.8 需要 22 个字，达到 0.9 需要 83 个字。

所以 top-10 保留的那 10 个字，加起来的概率大概在 0.6 到 0.7 之间。 剩下 0.3 到 0.4 的概率被直接丢掉，然后强行放大到 1。

这意味着：原本模型认为有 35% 概率”不会写这 10 个字里的任何一个”， 现在这个可能性被抹掉了。top-k 是一个相当激进的截断。

top-k 的致命问题：k 是一个固定的数，但分布的形状不是。

考虑两种极端情况：

• 分布很尖的时候（比如模型正在写一句它背得滚瓜烂熟的诗）： 前 3 个字就占了 99% 的概率。这时 k = 50 会把 47 个本该被排除的字 也放进候选集——引入了不必要的随机性。
• 分布很平的时候（比如模型正在决定一首诗的走向）： 前 50 个字加起来才 60% 的概率。这时 k = 50 砍掉的 是 40% 的正当可能性——又开始念咒了。

同一个 k，在一种情况下太宽，在另一种情况下太窄。 这是 top-k 的固有缺陷。

3.5.4 top-p：让累积概率决定候选集

top-p（也叫核采样，nucleus sampling）的思路是：不用固定的个数， 用固定的概率质量。

把所有候选按概率从大到小排序，从头开始累加，累加到刚好超过 p 为止：

S_p = \left\{ v : \sum_{u:\, p_u > p_v} p_u < p \right\} \tag{3.20}

所以这个式子在说：把字按概率排队，从最大的开始数， 数到累计概率超过 p 就停——这些就是候选集。剩下的丢掉，重新归一化。

p = 0.9 的意思是”保留能覆盖 90% 概率的那一小撮字”。

关键区别：S_p 的大小是浮动的。

• 分布尖的时候，可能 2 个字就覆盖了 90% → S_p 只有 2 个元素
• 分布平的时候，可能要 200 个字才覆盖 90% → S_p 有 200 个元素

候选集自动适应了分布的形状。 这正是 top-k 缺的那个东西。

实测的核曲线（τ = 1.0，词表 1200）：

累积概率阈值 p	保留字数	占词表比例
0.50	2	0.17%
0.80	22	1.8%
0.90	83	6.9%
0.95	185	15.4%
0.99	489	40.8%
0.999	889	74.1%

读这张表

最上面一行最惊人：p = 0.5 时只保留 2 个字。

也就是说，模型在”春江潮水连海平”之后，一半的概率质量压在前两个字上。 这符合我们的直觉——诗的第一句之后，接下来该写什么，模型其实相当有主意。

但注意最下面几行：p = 0.99 时要保留 489 个字， 占了词表的 40.8%。这说明这个分布的尾巴非常长—— 剩下 1% 的概率，稀稀拉拉摊在 489 个字上。

这个尾巴正是”纯采样”会翻车的地方：每个字概率都很低，但它们加起来有 1%。 生成 50 个字，期望会有 50 \times 1\% \times（尾巴上的字数权重）…… 虽然不是每步都中，但几十步里中一次就够跑偏了。

p = 0.9 保留了 83 个字（6.9%），这是实践中常用的默认值—— 既砍掉了长尾，又留够了多样性。

注记停下来想一想 ⑰

top-k 和 top-p 都在做”截断 + 重新归一化”。区别只在怎么决定截在哪里。

现在考虑：如果模型输出的分布是完全均匀的（1200 个字每个概率都是 1/1200）， top-p = 0.9 会保留多少个字？top-k = 10 呢？

（答案：均匀分布下，累积到 0.9 需要 0.9 \times 1200 = 1080 个字。 所以 top-p = 0.9 保留了 1080 个——几乎等于不截断。 而 top-k = 10 仍然只保留 10 个，强力地把分布变尖了。

这说明：top-p 在模型”没主意”的时候会退让，top-k 不会。 这也是为什么现在的主流做法是用 top-p 而不是 top-k。）

3.5.5 顺序陷阱：先截断还是先调温

到了这里，一个自然的问题：温度、top-k、top-p 可以一起用吗？

可以，而且实践中经常一起用。但这里有一个真实的坑：顺序不同，结果不同。

两种写法：

写法 A（先截断后调温）： 1. 用原始分布决定 top-p 的截断位置 2. 保留那些字，重新归一化 3. 在保留集内，用温度 τ 重新标定

写法 B（先调温后截断）： 1. 先用温度 τ 把整个分布重新标定一遍 2. 在调过温的分布上决定 top-p 的截断位置 3. 保留、重新归一化

τ = 1 时两者完全一样（除以 1 什么也没改变）。但 τ ≠ 1 时，候选集的大小会差很多。

实测（top-p = 0.9）：

τ	写法 A 保留	写法 B 保留	差多少倍
0.5	83	2	41.5 倍
1.0	83	83	1 倍（相同）
1.5	83	454	5.5 倍
2.0	83	673	8.1 倍

这个陷阱有多隐蔽

写法 A 的保留数永远是 83，不随 τ 变。这不是巧合—— 因为写法 A 的截断只看原始分布，跟 τ 无关。

写法 B 的保留数随 τ 剧烈变化： - τ = 0.5（让分布变尖）→ 保留 2 个。分布尖了，两个字就够 90% 了。 - τ = 2.0（让分布变平）→ 保留 673 个。分布平了，要 673 个字才够 90%。

两种写法不可能都对，它们描述的是不同的采样分布。

哪个”对”？没有标准答案——这取决于你想让 τ 和 top-p 各自控制什么。 但有两个后果必须知道：

1. 同一个 (τ, p) 参数组合，换一个库可能行为不同。 HuggingFace transformers 的 generate() 在同时给 temperature 和 top_p 时， 走的是先调温后截断。而不少自研实现写的是先截断后调温。 把参数从一处搬到另一处，输出会变。
2. 在写法 B 下，τ 和 p 是耦合的。 你以为调 τ 是在调”随机程度”， 实际上它同时改变了候选集的大小——两个旋钮在拧同一个东西。

我们自己的实验脚本 run_ch3_case.py 里，sample_next() 的 order 参数 就是用来切换这两种写法的。正文里所有 top-p 的实验，都用写法 A 报告。

提示怎么避免这个坑

只用一个旋钮。 实践中更稳的做法是：

• 用 top-p 控制候选集（它自适应）
• 用 温度控制候选集内的随机程度
• 两者只调一个，另一个固定

如果两个都要调，就在文档里明确写清楚顺序，并且测一遍极端值： 把 τ 拉到 0.5 和 2.0，看保留集大小是不是你预期的。上面的表就是干这个用的。

3.5.6 为什么不用 beam search

到这里，一个自然的想法是：既然后果这么严重，那我们把搜索做得更好一点不行吗？

贪心只看一步。那就看多步——beam search（束搜索）： 每一步保留 k 条最优的部分序列，最后选整体概率最高的那条。

\text{score}(y_{1:t}) = \sum_{i=1}^{t} \log p_\theta(y_i \mid y_{1:i-1}) \tag{3.21}

所以这个式子在说：一条候选序列的得分，是它每一步对数概率的累加。 因为概率相乘在对数下变成相加。得分越高（越接近 0），序列越”顺理成章”。

beam search 的目的是最大化这个得分——找一条整体概率最高的序列。 听起来完全正确。但它也失败，而且失败得更隐蔽。

beam search 的两个问题

问题一：最大化序列概率 ≠ 生成好文本。

这是 Holtzman 等人 2020 年那篇论文的核心论断。 考虑两个候选：

候选	每步概率	序列得分
A：春风又绿江南岸	0.3, 0.4, 0.35, 0.5, 0.3, 0.4, 0.45	\log 和 = -8.9
B：春风春风春风春风	0.5, 0.9, 0.85, 0.9, 0.88, 0.9, 0.86	\log 和 = -2.3

B 的得分远高于 A。 因为重复的内容模型预测起来极有把握—— 这正是 3.7 节会定量分析的”模型在退化文本上更自信”。

所以 beam search 会主动选择 B。它不是”搜索得不够充分”， 而是它最大化的那个目标本身就是错的。

问题二：beam search 输出的是”平均脸”。

因为它在最大化概率，而概率高意味着”常见”。 一篇全是常见句子的文本，读起来就是没有特色—— 每个词都最保险，合起来最平庸。

这和 3.7 节说的”用困惑度做筛选会偏好平庸输出”是同一个机制。

那 beam search 什么时候是对的

在有标准答案的任务上。 机器翻译、语音识别、代码补全—— 这些任务里”最可能的输出”通常就是”正确的输出”， 因为输出空间被输入严格约束了。

在开放生成任务上它是错的。 写诗、写故事、聊天—— 这些任务里没有唯一正确答案，“最可能”等于”最没意思”。

同一个算法，任务不同，结论相反。 这和 3.7.3 节 “置信度在有的任务上有用、有的任务上有害”是同一件事的两种表现。

注记停下来想一想 ⑱

本章的实验中，我们没跑 beam search。为什么？

（答案：两个原因。

工程上：beam search 要把 k 条候选序列都维护在缓存里， KV 缓存要按 beam 复制。本案例的模型太小，测出来的差异会被别的噪声淹没。

教学上：beam search 的失败机制和贪心是同一个—— 都是”最大化概率”，只是一个看一步、一个看多步。 文献里已经把这条路走完了（Holtzman 2020 的第 3 节）， 我们没必要重复一个已经有确定结论的实验。

知道什么时候不该做实验，和知道怎么做实验一样重要。）

3.5.7 回到案例：退化消失了

现在把 3.1 节的表再拿出来看一遍，这次我们知道每一列是怎么来的了。

解码策略	distinct-2	最长重复	为什么
贪心	0.102	8-gram	确定性系统，必然进入周期
贪心 + 屏蔽 unk	0.531	4-gram	换了个周期，性质不变
τ = 0.8	1.000	无	引入随机性，打破周期
top-k = 10	0.980	无	同上，但候选集固定
top-p = 0.9	1.000	无	候选集自适应
top-p = 0.9, τ = 1.2	1.000	无	两者叠加

循环消失了，因为 argmax 那个”必然”没了。

再看一眼 τ = 0.8 的输出：

春江潮水海平，狂唱山城秋。作城外由流風日裏，北樓殘月在晴雲。
葉春無一日來，攀折綠江中。清誰與一留人，一來少年春。

和贪心的春江潮水海平。，。，。，对比，这是同一套权重生成的。 模型没变，变的是”怎么从分布里挑一个字”这个动作。

这一节真正要你记住的

模型的能力和模型的输出是两件事。

前者由权重决定，后者由权重加上解码策略共同决定。 一个 loss 4.141 的模型，配上贪心解码，输出可以差到 distinct-2 = 0.102； 配上 τ = 0.8，输出可以到 distinct-2 = 1.000。

中间那 0.9 的差距，不是训练出来的，是解码策略给的。

这是本章”两副面孔”的第一层含义：训练决定模型知道什么， 解码决定模型说什么。两者独立，都必须懂。

3.5.8 怎么量”有多重复”：distinct-n 及其局限

3.1.3 节用 distinct-n 判断退化程度。把它的定义写清楚， 因为一个你不理解的指标，会误导你。

定义：把文本切成所有长度为 n 的连续片段， 不重复的片段数 ÷ 总片段数。

\text{distinct-}n = \frac{|\{\text{所有长度为 } n \text{ 的片段}\}|}{T - n + 1} \tag{3.22}

所以这个式子在说：如果每个二字组合都不重复，distinct-2 = 1； 如果一半是重复的，distinct-2 = 0.5。

用本章的贪心输出走一遍。文本是（<unk> 用 U 代替）：

春 江 潮 水 海 平 U 。 ， U 。 ， U 。 ， U …

distinct-1：不重复的单字有几个？ 春 江 潮 水 海 平 U 。 ，——9 个。总长度 57。 9/57 = 0.158。实测表里是 0.060——差异来自输出的具体长度和 decode 的处理， 但量级一致：不到 16% 的字是”新”的。

distinct-2：二字组合。前 10 个是 春江, 江潮, 潮水, 水海, 海平, 平U, U。, 。，, ，U —— 9 个里没有重复。 但从第 10 个开始全是 U。 和 。， 的交替，只有 2 个不同的组合。 所以 11/56 \approx 0.196。实测 0.102，同一个量级。

distinct-n 的三个局限

局限一：它惩罚合理的重复。

真实文本里也有重复：「青青河畔草」的「青青」是刻意重复， 「尋尋覓覓，冷冷清清，淒淒慘慘戚戚」全是重复—— 它们的 distinct-2 会很低，但它们是好文本。

所以 distinct-n 只能用来发现异常，不能用来评价质量。

局限二：它对长度敏感。

生成 10 个 token 时，distinct-2 一定很高（因为片段少，撞不上）。 生成 10000 个 token 时，重复几乎是必然的。 不同长度的文本之间不能直接比 distinct-n。

局限三：它抓不住”语义重复”。

「不見君不見，不見，不如何處」——distinct-2 可能不低 （因为「不見」后面跟的字在变），但读起来就是复读。

这是 n-gram 类指标的共同缺陷：它们只看表面的片段匹配， 不看语义。这也是为什么评估生成质量至今没有好指标—— BLEU、ROUGE、perplexity 各有各的失效场景。

本章的用法是安全的：我们只在同一个模型、同样长度、 同样提示词下横向比较不同解码策略。这三个因素都控制住了， distinct-n 的局限就不会干扰结论。

注记停下来想一想 ⑲

3.1.3 节的表里，top-k = 10 的 distinct-2 是 0.980， 是四种采样策略里最低的。

而它的 top-k 只保留了 10 个候选——比 top-p = 0.9 的 83 个少得多。

候选集更小，按理说应该更”保守”、更容易重复。 数据符合这个直觉。

但 τ = 0.8 的候选集是全部 1200 个（温度不截断）， 它的 distinct-2 是 1.000——比 top-k = 10 更不重复。

所以”候选集大小”和”重复程度”不是单调关系？

（答案：不是。因为重复的来源不是候选集大小，而是”确定性”。

top-k = 10 保留 10 个候选，但如果这 10 个里有一个概率是 0.95， 那它几乎每次都抽中同一个——行为接近贪心。

而 τ = 0.8 不截断，但整个分布被压平了， 最大的那个候选概率只有 0.632（见 3.5.2 节表）， 所以抽中同一个字的概率低得多。

决定重复程度的是”分布的尖锐程度”，不是”尾巴有多长”。 top-k 砍尾巴但不动尖峰；温度直接压尖峰。这就是两者效果不同的原因。）

3.5.9 真实系统怎么配这些参数

本节把理论落回实践。真实推理服务的默认参数是怎么定的， 以及为什么不同的服务默认值不一样。

常见配置（2024–2025 年的主流做法）：

场景	温度	top-p	说明
代码生成	0.2 – 0.4	0.95	要确定性，不允许多样
数学推理	0.0 – 0.3	1.0	同上，且通常配 best-of-n
通用对话	0.7 – 1.0	0.9 – 0.95	平衡
创意写作	0.9 – 1.2	0.95 – 1.0	要多样性
翻译	0.0 – 0.3	1.0	有标准答案

这张表背后的原则

看任务有没有”正确答案”。

第 3.5.6 节讲过：beam search 和有标准答案的任务是匹配的， 因为它们都在最大化”最可能”。温度也一样：

• 有标准答案 → 用低温度，让模型给出最可能的答案。 代码、数学、翻译都属于这类。
• 没有标准答案 → 用高温度，让模型探索。 写诗、起名字、头脑风暴都属于这类。

但注意”低温度”不等于”温度 0”。 代码生成常用 0.2 而不是 0， 因为完全确定的输出会反复犯同一个错误—— 同一个 bug 你重试十次，得到十次一样的错误代码。 保留一点随机性，至少给你一个换个思路的机会。

这是一个纯粹的工程考量，不是数学结论。

警告两个真实系统里的坑

坑一：参数名相同，语义不同。

本章 3.5.5 节的顺序陷阱不是理论问题。同一个 (temperature, top_p)， 在不同的库里可能有不同的行为：

• 有的先调温后截断（HuggingFace 的部分版本）
• 有的先截断后调温
• 有的还有 top_k 同时生效，且顺序又不一样

把提示词和参数从一个服务搬到另一个服务时，输出会变。 这不一定是”新服务更差”——可能只是顺序不同。

坑二：temperature=0 的实现不统一。

有的库在 temperature=0 时报错（除零）； 有的把它特判成 argmax； 有的把它当成一个极小数（比如 1e-6）然后照常采样。

最后一种最危险：它看起来是贪心，但在概率接近的两个 token 上 可能给出不同结果。跨平台复现实验结果时，这会成为无法解释的差异来源。

做法：如果你的代码依赖”确定性输出”，不要依赖 temperature=0， 直接调用 argmax。

注记停下来想一想 ⑳

本章的案例脚本里，贪心解码用的是 argmax，而不是 temperature=1e-6。

但如果用 temperature=1e-6，结果会一样吗？

（答案：在数学上是极限相同，在浮点上是几乎相同，但不保证逐位相同。

\tau = 10^{-6} 时，logits 被放大 10^6 倍。 如果两个候选的 logits 差值是 10^{-7}——在 float32 里这是可表示的—— 放大后差 0.1，softmax 后概率比约 e^{0.1} = 1.105。 不是 argmax。

所以”温度趋近 0 就是 argmax”是一个数学极限， 在有限精度的浮点数上不成立。

这又是一个”数学上等价 ≠ 实现上等价”的例子， 和 3.9.5 节那个”K/V 都倒序”是同一类问题。）

3.6 上下文长度：一条硬边界

3.6.1 64 这个数字是怎么来的

本模型的 block_size = 64。这个数字不是随便定的，它决定了三件事：

1. 位置嵌入表的大小：pos_emb 是 (64, 128)，只有 64 行。 第 65 个位置没有对应的行可以查——这就是 3.9.6 节那个 CUDA 崩溃的根源。
2. 因果掩码矩阵的大小：(1, 1, 64, 64)。它在模型构造时就固定下来了。
3. 训练时一次前向能看到的最长文本：get_batch 每次从语料里切 64 个连续字符。

三个限制里，第一个是硬的，后两个是软的。

掩码矩阵可以重建成更大的（换个 block_size 重新初始化即可）； 训练样本的长度可以改。但位置嵌入表只能有 64 行， 因为它是在训练中被学出来的——第 65 行从来没被训练过， 你就算把它加上去，里面的数值也是随机的。

“训练时的上下文”和”能用的上下文”是两回事

这三个限制指向同一个事实：模型的能力边界由训练时见过什么决定，不是由代码允许什么决定。

你可以把 block_size 改成 4096，让代码”支持”长上下文。 但 pos_emb 的第 65 到 4096 行是随机初始化的—— 模型在第 100 个位置上的行为，等同于在一张白纸上乱猜。

这就是为什么”扩展上下文长度”从来不是改一个数字那么简单。 它要么需要重新训练（贵），要么需要一套位置编码的外推方案 （RoPE 插值、NTK 缩放等，第 4 章会讲）。

第 2 章我们测过 RoPE 的外推——结果是没测出预期的效果， 因为那个任务根本不需要长程依赖。这个教训在这里同样适用： 上下文长度是一个能力上界，不是需求。

3.6.2 为什么注意力是 O(T²)

第 3.2.4 节提过，注意力的计算量是 O(T^2 d)，前馈是 O(T d^2)。 T 小的时候前馈是大头，T 大的时候注意力反超。

算一下交叉点。令：

T^2 d = T d^2 \quad \Longrightarrow \quad T = d \tag{3.23}

所以这个式子在说：当上下文长度等于模型维度时，两项持平。

本模型 d_{model} = 128，T = 64——注意力项还是前馈项的一半， 所以注意力不是瓶颈。

Llama-3-70B：d_{model} = 8192，训练上下文 8192。正好在交叉点上。 当上下文扩到 128K 时，T/d = 16，注意力项是前馈项的 16 倍—— 它成了绝对瓶颈。这就是为什么长上下文模型都在研究稀疏注意力、 滑窗注意力、线性注意力这些降低 T^2 项的方案。

一个具体的数字

T = 128000 时，注意力矩阵有 T^2 = 1.64 \times 10^{10} 个元素。

用 float16 存，一个元素 2 字节，光是这一个矩阵就要 32.8 GB。

注意这是一层、一个头的量。70B 模型有 80 层、64 个头。 当然实际实现不会真的物化这么大的矩阵（会分块计算）， 但计算量本身没有变：1.64 \times 10^{10} 次乘加， 即使按 100% 利用率算，也要跑很久。

这就是为什么”长上下文”是一个工程难题，而不是一个开关。

3.6.3 超出边界时会发生什么：三种表现

回到 3.9.6 节那个崩溃。把”超出上下文”这件事的三种表现列清楚， 因为它们的排查难度完全不同。

表现一：CUDA device-side assert（我们遇到的）

indexSelectSmallIndex: block: [0,0,0], thread: [64,0,0]
Assertion `srcIndex < srcSelectDimSize` failed.

最难查的一种。报错信息里没有”位置”“上下文”“block_size”任何一个词， 而且它会污染 CUDA 上下文——后续所有 CUDA 操作都失败， 必须重启进程才能继续。

表现二：静默截断

有些实现会把超出部分悄悄丢掉：

x = x[:, -block_size:]        # 只看最后 block_size 个

程序不报错，但模型看到的历史和你以为的不一样。这种最难发现—— 输出”看起来还行”，只是质量莫名下降。

表现三：位置回绕

如果实现用取模来避免越界：

pos = torch.arange(T) % block_size

第 65 个位置会变成位置 0。模型会以为自己在写句首—— 这和 3.9.2 节那个 pos_zero bug 是同一种错误， 只不过它是”每 64 个字符发生一次”。

怎么区分这三种

看 logits 的形状：形状正常 → 不是表现一。

看位置索引：打印 pos 的最后一个值。如果它不增长（一直是 0-63 循环）， 那是表现三；如果增长但超过了 block_size，那是表现一或二。

看生成结果的周期性：表现三会产生每 64 个 token 一次的规律性抖动—— 因为模型每 64 步就”以为”自己重新开头了。 把生成结果按 64 分组看，如果每组的开头质量明显不同，基本可以确诊。

本章的 forward_step() 用的是第四种处理方式：提前报错。 它在越界之前抛一个说人话的 ValueError， 把”CUDA device-side assert”翻译成”位置嵌入表只有 64 行”。

这是工程上最值得推荐的做法：与其让错误在最难懂的地方暴露， 不如在你能解释清楚的地方拦住它。

注记停下来想一想 ㉑

上面三种表现，哪一种在单元测试里最容易漏掉？

（答案：表现二，静默截断。

表现一会让程序崩溃，测试必然失败； 表现三会让输出出现周期性异常，如果测试里有质量检查也可能被抓到。 但表现二什么都不做——程序正常跑，输出正常返回， 只是”悄悄变得差一点”。

静默降级是最危险的失败模式，因为它不触发任何告警。 测试里发现不了，只能靠监控——比如监控”输入长度分布”和 “实际被模型看到的长度分布”之间的差异。）

3.6.4 滑动窗口：一个不免费的替代方案

既然硬边界会崩，那就不设边界——只保留最近 W 个 token 的 K/V。

def append(self, layer, k, v, window=None):
    self.k[layer] = torch.cat([self.k[layer], k], dim=-2)
    self.v[layer] = torch.cat([self.v[layer], v], dim=-2)
    if window is not None and self.length > window:
        self.k[layer] = self.k[layer][:, :, -window:, :]     # 丢掉最旧的
        self.v[layer] = self.v[layer][:, :, -window:, :]
    return self.k[layer], self.v[layer]

它解决了显存和计算量问题：缓存大小恒定为 W， 注意力计算量也从 O(T^2) 降到 O(T \cdot W)。

但它引入了两个新问题：

问题一：模型永远只看得到最近 W 个 token。 这是能力上的妥协，不是免费的。如果任务需要记住 1000 个字之前的信息， 滑动窗口直接把这个能力砍掉了。

问题二：位置怎么编号。 - 如果继续用绝对位置（W 之后的位置编号继续增长）， 则位置编号会超过 pos_emb 的行数——又回到了硬边界。 - 如果位置重新从 0 编号，则每个窗口的开头都被当成句首—— 模型会以为每 W 个 token 就是一整句话。

两个选择都是错的，这就是绝对位置编码的死结。 第 2 章讲 RoPE 时提到的”相对位置”就是解法： RoPE 编码的是两个位置之间的差，窗口滑动不改变任何一对 token 的相对距离， 所以位置编号怎么变都不影响。

提示一个真实系统怎么做

现代长上下文模型的组合拳通常是三层：

1. RoPE（位置编码天然相对，不惧滑动）
2. 滑动窗口 + 全局 token：大部分层用窗口注意力（便宜）， 少数层用全局注意力（贵但能记住远处）。Gemma、Mistral 用这个思路。
3. attention sink：在窗口最前面永远保留第一个 token。 实证发现模型会把大量注意力分配给第一个 token—— 它像一个”什么都不关注时的默认去处”。把它丢掉会让模型行为异常。

第 3 点是纯经验发现，理论解释至今没有共识。 它是”实践跑在理论前面”的典型例子——这类发现在这个领域很常见。

这一节真正要你记住的是

上下文长度是一个三维约束，不是一维的。

维度	限制来自	能不能扩
位置嵌入	训练时见过多少位置	需要重训或外推方案
注意力计算量	O(T^2)	需要算法改进（稀疏/线性）
KV 缓存显存	O(T)、随生成增长	需要窗口或分页管理

三个维度必须一起解决。只改一个数字（比如把 block_size 调大） 只会让程序跑起来，不会让模型变强——这中间的差距， 就是”能动”和”能对”的差距。

3.7 训练与推理的错配：一个错 token 的阴影

3.7.1 teacher forcing 是一份”完美前文”的礼物

回到 3.2.5 节那个观察。

训练时，模型在位置 t 看到的输入是 x_1, \dots, x_t——真实的文本。 这在文献里叫 teacher forcing（教师强制）：每一步都把正确答案喂给它， 不管它上一步预测得对不对。

这是个高效的训练技巧——它让整个序列可以并行计算，是 3.2 节全部推导的前提。 但它有一个副作用：

模型从来没有练习过”在错误的前文下继续写”。

用图看会很清楚。假设真实文本是「白日依山尽」：

训练时（teacher forcing）：
  输入 白 →  预测 日  ✓
  输入 白 日 →  预测 依  ✓
  输入 白 日 依 →  预测 山  ✓
  输入 白 日 依 山 →  预测 尽  ✓
  每一步的输入都是真实的（蓝色）

推理时（自回归）：
  输入 白 →  预测 日  ✓
  输入 白 日 →  预测 落  ✗   ← 从这里开始，前文就是假的了
  输入 白 日 落 →  预测 ，  ✗
  输入 白 日 落 ， →  预测 。  ✗   ← 越走越远

到了第 3 步，模型面对的输入「白日落」是训练时几乎不可能出现的句子。 它在这个分布外的输入上的行为，没有任何保证。

为什么叫”分布外”

训练集是全唐诗。里面「白日」后面跟「依山尽」的概率很高， 跟「落」的概率很低——但不是零。「白日落」这种组合在诗里是存在的。

真正的问题是组合爆炸。模型训练时见过的所有 64 字片段， 在整个可能片段空间里是极小的一部分。一旦生成过程中偏出这个集合， 模型就进入了它没见过的区域。

而生成过程一定会偏出去——因为只要有一处 argmax 和真实文本不同， 后面的输入就再也不是训练时的分布了。

3.7.2 实测：一个错 token 的阴影有多长

“错配”听起来很抽象。我们把它量化。

实验设计：

1. 取一段真实的验证集文本，让模型 teacher forcing 一遍， 记下每个位置的 loss，作为基线。
2. 挑一个位置 k，把模型”想要写的字”（argmax，排除 <unk>）替换掉 真实的第 k+1 个字——只改这一个。
3. 从 k+2 开始重新喂真实文本，再算一遍每个位置的 loss。
4. 两条曲线的差 \Delta \text{loss}(t)，就是”位置 k 那个错误”在 位置 t 留下的阴影。

实测结果：

位置 k	真值	模型想写的	下一位置 \Deltaloss	峰值	尾部均值	回落一半用了
4	方	日	+2.485	+2.485	+0.126	3 步
8	御	，	+10.883	+10.883	+0.388	2 步
16	<unk>	不	+0.621	+0.621	−0.013	2 步
24	蓋	首	−0.219	+0.231	−0.015	1 步

这张表读出来的三件事

第一，错误的代价极其悬殊。 位置 8 的替换让下一位置的 loss 暴涨 10.883， 而位置 24 的替换反而让 loss 下降了 0.219（模型的错误猜测恰好碰对了上下文）。

同一个模型、同一个文本、同样”改一个字”，代价从 −0.22 到 +10.88， 跨了两个数量级。 这说明”错配有多严重”这个问题没有统一答案—— 它取决于你错在哪里。

第二，阴影衰减得很快。 四个案例的”回落一半”分别用了 3、2、2、1 步。 位置 16 和 24 的尾部均值甚至是负的（−0.013、−0.015）—— 意思是错误造成的损害在十来步之后已经完全消失。

这个结果有点出乎意料。文献里常说的”exposure bias 导致误差累积”， 在这个任务上不明显。

第三，为什么？ 因为这是一个字符级的唐诗任务，上下文长度只有 64， 而且第 2 章已经测过：这个任务对长程依赖的需求很弱 （上下文从 8 加到 128 只改善了 0.64，而且 16 比 8 还差）。 在弱依赖的任务上，错误的传播自然很快被冲淡。

结论要按任务说，不能按论文说。 这是本书反复强调的一点。

警告那循环是怎么来的

如果错误的阴影只有 1–3 步，为什么贪心解码会陷入长达 8 个字的循环？

因为循环不是”一个错误的持续影响”，而是”错误在持续产生”。

区别在这里： - 阴影衰减问的是：给定一个错误，它影响多远？ - 循环问的是：每一步都在产生新的错误，合起来会怎样？

贪心解码下，每一步都选 argmax。argmax 在退化文本上是自洽的—— 一旦输出「。」，下一个位置 argmax 大概率还是「，」， 因为「。，」这个组合在语料里出现过太多次，模型对它极度自信。

所以循环是一个稳定不动点，不是误差累积的结果。 3.3.3 节的确定性动力系统分析才是正确的解释框架。

3.7.3 置信度错位：模型什么时候最自信

再看一组数。这一组是本章最值得玩味的。

在真实文本上，模型给自己的 argmax 的平均概率：0.239（平均 NLL 2.203）。

在自由生成时，模型对自己所选 token 的平均概率：0.157（平均 NLL 2.135）。

换个角度看这组数

表面上看，自由生成时模型”更不自信”（0.157 < 0.239）。 但换成 NLL 看，自由生成时它反而更自信（2.135 < 2.203）。

这个矛盾是因为平均的对象不一样： - 0.239 是”在真实文本上，模型押的那个字的概率”的平均 - 0.157 是”生成时，模型押的那个字的概率”的平均

两者的概率分布形状不同，直接比较均值会误导。

正确的问题是：模型在什么情况下更”笃定”？

答案：在它自己生成的、退化的文本上更笃定。 因为它生成的文本高度重复，「。，」这种组合模型见过无数次， 预测起来自然容易（NLL 更低）。

而这份”笃定”恰恰是循环的燃料。 越笃定 → 越容易 argmax 到同一个字 → 前文越重复 → 下一步越笃定。这是个正反馈。

一个危险的推论

如果模型在退化文本上更自信，那么“用模型的置信度做筛选”是有害的。

实践中有人这么做：生成多个候选，挑模型自己概率最高的那个。 或者：在生成时监控置信度，置信度低就重新生成。

在这个模型上，这套逻辑会主动挑出最退化的那个输出。

更一般地说：模型的置信度衡量的是”这句话像不像训练数据”， 不是”这句话对不对”。 训练数据里有大量重复（诗歌里有叠字、有复沓）， 所以重复的文本在模型眼里是”正常”的。

这个陷阱在真实系统里同样存在。一个用困惑度做筛选的重排序器， 会偏好那些”平庸但熟悉”的输出——这正是 Holtzman 等人 2020 年那篇论文 要解决的问题。

注记停下来想一想 ㉒

上面说”模型的置信度衡量的是像不像训练数据，不是对不对”。

那么反过来问：有没有一种情况，模型非常自信，但输出确实是好的？

（答案：有，而且很常见。当任务简单、或者模型在复述训练数据里 高频出现的模式时——比如让它续写”床前明月光”， 它会极自信地写出”疑是地上霜”，而这确实是正确答案。

所以在有标准答案的任务上（翻译、摘要、代码补全）， 置信度是有用的信号；在开放生成任务上（写诗、写故事）， 高置信度反而可能是退化的标志。同一个指标，任务不同，含义相反。）

3.7.4 一个自然的想法：让训练也看错误前文

既然问题出在”训练时永远看真实前文”，那训练时也掺一些模型自己的输出不就行了？

这个想法叫 scheduled sampling（计划采样，Bengio 等人 2015）。 做法是：训练的第 t 步，以概率 \epsilon_i 用模型自己的预测作为下一步的输入， 以概率 1 - \epsilon_i 用真实 token。\epsilon_i 随训练进度从 0 慢慢升到某个值。

早期（ε≈0）：完全 teacher forcing —— 先学会基本的语言能力
后期（ε>0）：逐渐掺入自己的错误 —— 学会从错误中恢复

听起来完美。但它有几个真实的困难，导致它没有被主流采用。

scheduled sampling 的三个问题

问题一：它破坏了并行训练。

3.2.3 节说过，训练能并行的前提是”位置 t 的输入只依赖位置 1..t 的真实 token”。 一旦下一步的输入可能是模型自己的预测，这个链条就依赖前一步的计算结果—— 必须串行。

换句话说：scheduled sampling 把训练也变成了推理的样子， 速度直接掉到 1/T。这是不可接受的代价。

问题二：目标和推理仍然不一致。

就算训练时掺了错误，推理时的错误分布和训练时掺进去的错误分布是不同的。 训练时模型还”半生不熟”，它犯的错和训练完成后犯的错不一样。 所以这是一个移动靶——你在用旧错误训练一个会产生新错误的模型。

问题三：理论上的不一致性。

Huszár 等人 2015 年指出，scheduled sampling 优化的是一个 和原始目标不同的目标。它在某些设定下会让模型学到 “根据前一个 token 的置信度来决定要不要相信它”这种捷径， 而不是真正的纠错能力。

所以它被更简单的方案取代了——见下。

现在主流的做法是什么

不训练纠错，训练鲁棒性。 具体来说是三件事：

1. 接受这个错配，把它当成”采样带来的多样性”。 既然采样本来就要引入随机性，那模型见到非典型前文就是正常的。
2. 用大规模数据覆盖。当训练数据有几十亿条不同上下文时， 模型很可能见过”类似”的错误前文——不是完全没见过。
3. 在推理侧用采样而不是贪心（3.5 节）。 采样产生的分布更接近训练分布，因为它不追求”唯一的正确答案”。

第 3 点最值得注意：它说明训练和推理的错配， 主要靠推理侧解决，而不是训练侧。这是本书”两副面孔”主题的又一次体现—— 你没法把两面合成一面，只能在两面之间找一个都能接受的工作点。

注记停下来想一想 ㉓

本章 3.7.2 节的实测显示：在这个任务上，一个错 token 的阴影 1-3 步就衰减一半。

如果这个结论对所有任务都成立，那”训练/推理错配”就不值得担心了。

它成立吗？

（答案：不成立。请回想第 2 章那两个实验：

上下文长度实验：8 → 128 改善了 0.64，说明这个任务对长程依赖需求弱。 RoPE 外推实验：各种位置编码都没有明显差异，同一个原因。

唐诗续写是一个”局部任务”——下一句诗主要靠上一句的意象和韵脚决定， 不依赖整首诗的结构。在这样的任务上，错误自然很快被冲淡。

换一个任务：写一个 500 行的程序。 第 3 行少了个括号，第 200 行会是什么样？错误的传播是全局的。

所以 3.7.2 节那个”阴影很短”的结论， 必须和”这是唐诗任务”这个前提一起引用。 脱离任务谈”错配严重不严重”是没有意义的。）

3.7.5 阴影的形状：它不单调衰减

3.7.2 节的表给了三个汇总数字（峰值、尾部均值、半衰步数）。 但只看汇总会错过最重要的东西：阴影的形状。

把位置 8 那个替换（真值「御」被换成「，」）之后的逐位置 \Deltaloss 画出来：

距扰动几步	\Deltaloss
+1	+10.883
+2	+1.056
+3	+2.536
+4	−0.486
+5	+0.246
+6	+2.366
+7	0.000
+8	0.000
+9	+0.252
+10	−0.323
+11	+0.260
+12	−0.115
+13	−0.095
+14	+0.461

这条曲线在说什么

它不是一条衰减曲线。 它在 +10.9 → +1.1 → +2.5 → −0.5 → +2.4 … 之间来回摆动，而且在 14 步之后还没有回到 0。

三个观察：

第一，“半衰 2 步”这个汇总量具有误导性。 从 +10.883 降到 +1.056（不到一半）确实只用了 1 步， 但第 3 步又反弹到 +2.536，第 6 步又到 +2.366。 说它”2 步衰减一半”是对的，但它没有继续衰减。

第二，振荡的原因是”诗歌的韵脚和结构”。 扰动让模型偏离了原诗，但诗的韵律结构还在拉它回来。 每次它快要回到正轨时，下一个字的格律要求又把 loss 推高。 这是任务结构造成的，不是模型的普适性质。

第三，尾部均值 +0.388 掩盖了 ±0.5 的振幅。 均值的绝对值很小，但每一项都在正负之间跳。 又是一个”均值误导”的例子——和 3.2.2 节那张表里 “平均 loss 5.024 掩盖了 0.000 到 8.184 的跨度”是同一类问题。

警告这对”误差累积”的教科书说法意味着什么

文献里常见的表述是”exposure bias 导致误差累积”—— 听起来像一个单调增长的过程。

实测出来的不是累积，是振荡。

区别很重要： - 累积意味着”时间越长越糟”，对策是”限制长度” - 振荡意味着”错误的即时影响很大，但会被任务结构吸收”， 对策是”保证每一步都别错得太离谱”

后者才是本章数据支持的结论。

但请记住 3.7.4 节那个”停下来想一想”的答案： 这是唐诗任务的性质，不是普遍规律。 在代码生成这类”错误会破坏全局结构”的任务上， 累积效应是真实存在的。引用结论时必须带上任务前提。

注记停下来想一想 ㉔

上面说这首诗的错误影响”被任务结构吸收”。

这个”结构”具体是什么？ 为什么它能把 loss 拉回来？

（答案：韵脚和平仄。

近体诗有严格的格律：偶句押韵、句内平仄交替。 这些约束极大地压缩了下一个字的可能范围—— 在押韵位置上，模型只需要在”同韵部的字”里挑，候选集本来就小。

所以当扰动让模型偏离后，下一步的格律约束会把它拽回来。

验证方法：把同一个实验在散文（不带格律的语料）上做一遍。 如果散文上的阴影不衰减（或者衰减慢得多）， 就说明”是格律在吸收错误”，而不是”模型有纠错能力”。

本章没有做这个对照实验——它是一个明确的后续工作。 把没做的实验说清楚，比假装做过要好。）

3.7.6 另一处错配：训练模式和推理模式

“训练与推理是两副面孔”还有一层和数学无关的含义： 同一个模型对象，在两种模式下行为不同。

PyTorch 里这是通过 model.train() / model.eval() 切换的。 本案例的模型 dropout = 0.0，所以两者没有区别—— 但真实模型几乎都会用 dropout（通常是 0.1）。

dropout 在做什么：训练时随机关掉一部分神经元（把它们置零）， 强迫模型不依赖任何单个神经元。推理时关掉 dropout， 用完整的网络。

为什么推理时必须关掉？ 因为 dropout 是随机的。 如果推理时还开着，同一个输入每次给不同的输出—— 这和 3.3.5 节讲的”贪心解码是确定性系统”直接冲突， 而且会让 3.4 节那个”缓存必须逐位相同”的测试永远失败。

忘记 model.eval() 是最经典的 bug 之一

它有几个特点，让它特别难查：

特点一：训练时完全看不出来。 train() 是默认状态， 你写训练循环时不需要显式调用它。所以”忘记切换”这件事 只会在推理时暴露。

特点二：症状是”输出不稳定”，不是”输出错误”。 同一个提示词生成两次，结果不同。 你很容易把它归因为”模型本来就随机”——如果用了温度采样， 这个解释甚至是对的。只有在贪心解码下才会发现异常。

特点三：它会污染缓存测试。 本章 3.4.9 节的测试一用 torch.equal 要求逐位相同。如果 dropout 还开着， 这个测试必然失败——而且差值会是 0.1 量级（和真 bug 一样大）， 完全无法用 3.4.3 节那张表区分。

排查方法：在测试函数的第一行加 model.eval()， 并在断言失败时先打印 model.training。

def test_cache_is_identity(model, ids):
    model.eval()                                    # ← 先加这一行
    assert not model.training, "model 还在训练模式"
    ...

这一个 assert 能省掉几个小时的排查。

还有哪些”模式相关”的行为

dropout 最常见，但不是唯一。PyTorch 里还有几个：

层	train() 时	eval() 时
Dropout	随机置零	恒等
BatchNorm	用当前 batch 的均值方差	用训练时累计的滑动平均
LayerNorm	无差别	无差别

注意第三行。 LayerNorm 在两种模式下行为相同—— 因为它的归一化只依赖当前样本自己的特征维， 不依赖 batch 内其他样本，也不需要累计统计量。

这正是第 2 章讲的 Pre-LN 能训深层模型的原因之一： 它没有 batch 相关的状态，所以不存在”训练和推理行为不一致”的问题。

BatchNorm 就不同了：它在两种模式下用的统计量不一样。 如果训练时的 batch 分布和推理时差很多 （比如推理时 batch = 1），BatchNorm 会引入一个真实的分布偏移。

这就是为什么现代大语言模型几乎不用 BatchNorm—— 序列长度可变、batch 可变，BatchNorm 的统计量根本不稳定。 LayerNorm 和 RMSNorm 是标准选择。

注记停下来想一想 ㉕

本章的模型 dropout = 0.0，所以 train() 和 eval() 没有区别。

为什么本章要特意设成 0？

（答案：为了让”缓存是恒等变换”这个命题在测试里成立。

如果 dropout = 0.1，那么即使正确调用 model.eval()， 前向传播仍然有 10% 的随机性——但 eval() 会关掉它， 所以理论上测试仍然能过。

问题在于如果哪一步忘了 eval()，测试就会失败， 而失败的原因和缓存实现毫无关系。这会给排查引入噪声。

更好的理由：本章的模型只有 42 万参数，训练 2000 步， 根本不会过拟合，dropout 没有收益。

第 2 章讲过 dropout 的作用是正则化—— 在小数据、大模型时才有意义。本章两者都不满足。

这个选择是”消除无关变量”，不是”dropout 不好”。

3.7.7 训练目标和推理目标本来就不是一回事

本章讲到这里，可以回答最初那个问题了。但在回答之前， 先把一个更根本的事实摆出来——它是本章所有现象的共同根源。

训练在做什么：最小化交叉熵

\min_\theta \; \mathbb{E}_{x \sim \text{数据}}\left[-\log p_\theta(x_{t+1} \mid x_{1:t})\right] \tag{3.24}

交叉熵是一个恰当评分规则（proper scoring rule）： 它在 p_\theta 等于真实条件分布 p_{\text{data}} 时取到最小值。 也就是说，训练的目标是让模型输出真实的概率分布。 它要的是整个分布对，不是某一个 token 对。

推理在做什么：从分布里选一个 token。

x_{t+1} = \text{select}\big(p_\theta(\cdot \mid x_{1:t})\big) \tag{3.25}

注意 select 这个函数在训练目标里根本不存在。

这一步之差，是本章全部现象的来源

把两件事并排看：

	训练	推理
输入	完整分布 p_\theta 和正确答案	只有分布 p_\theta
输出	一个标量（loss）	一个 token
目标	让分布接近真实分布	？

推理那一列的”目标”是空的。 这不是笔误—— 语言模型的训练目标里，从来没有定义过”怎么从分布里选一个 token”这件事。

所以 select 函数（贪心、温度、top-k、top-p）不是从训练目标推导出来的， 而是外加上去的。它们的合理性必须单独论证。

本章做的就是这个论证，而结论是：

• 贪心：最大化单步概率。它有一个漂亮的性质（确定性）， 但 3.3.5 节证明了确定性 + 有限状态 = 必然循环。
• 纯采样：无偏地从分布里抽。理论上”最忠实”， 但 3.5.1 节的实测显示最大概率只有 0.45， 长尾加起来有 1%——几十步里必然踩到一次。
• top-p：在”忠实”和”不退化”之间找一个工作点。 它没有理论保证，是实验调出来的（Holtzman 2020）。

没有一个是”正确”的。它们都是工程选择。

一个必须承认的结论

上面这段话等于承认：语言模型的生成质量，有一部分不由模型决定。

同一个模型、同一个提示词：

解码	distinct-2	输出
贪心	0.102	春江潮水海平。，。，。，
τ = 0.8	1.000	春江潮水海平，狂唱山城秋…

中间那 0.9 的差距，不在权重里，在解码器里。

这有好的一面：你不需要重新训练模型，就能大幅改变输出。 也有坏的一面：“这个模型效果不好”这个判断，可能根本不是在说模型。

所以评估一个生成系统时，第一件要问的是”解码参数是什么”， 而不是”模型多大”。本章 3.5.9 节那张配置表就是这个意思—— 它列的不是”哪个模型好”，而是”哪类任务该配什么解码”。

提示这一节和整本书的关系

上篇（第 1–4 章）讲的是”模型算什么”： 注意力怎么加权、Block 怎么组织、位置怎么编码。 这些都在决定 p_\theta 的形状。

本章第一次把话题从 p_\theta 移到了”怎么用 p_\theta“。

而下篇（第 5–8 章）会走得更远：扩散模型干脆放弃”选一个 token”这个框架， 改成”迭代地细化整个序列”。3.3.7 节已经点了一下这个方向。

换句话说：本章发现的那个”目标和手段之间的缺口”， 下篇给出的答案是——换一个目标。

那个答案是否更好，是下篇要回答的问题。

注记停下来想一想 ㉖

上面说”推理那一列的目标是空的”。

但真实系统明明有目标——比如”回答要准确”“要有帮助”。 为什么说它是空的？

（答案：因为训练目标里没有它们。

“准确”“有帮助”是人类偏好，不是交叉熵。 交叉熵只保证”模型输出的分布接近训练数据的分布”—— 而训练数据里既有好回答也有坏回答，甚至有很多平庸的回答。

把人类偏好变成可优化的目标，是另一个完整的研究方向 （RLHF、DPO、宪法 AI 等），本书不展开。

但有一点和本章直接相关：这些方法最终也要落到”怎么选 token”上。 PPO 训练出来的策略网络，推理时仍然要用温度、top-p 来采样。 本章讲的所有解码陷阱，一个都不会因为做了 RLHF 而消失。）

3.8 交互实验

本节三个实验，都是独立 HTML 文件，可以单独打开、全屏投影。 网页里直接嵌在下面，PDF 版对应位置是静态快照。

3.8.1 实验 3-1　采样实验室

这个实验用真实的模型分布（就是本章那个 42 万参数模型在 “春江潮水连海平”之后的输出，1200 维，导出成 JSON）。 你可以拖三个滑块：温度 τ、top-k、top-p，实时看到：

• 分布形状：1200 根柱子，被截断的部分变灰
• 保留集大小：多少个字还活着
• 熵：当前分布的困惑度
• 抽样结果：点一次抽一个字，看它从哪根柱子上掉下来

动手：先把 τ 固定在 1.0，把 top-p 从 1.0 慢慢拉到 0.5。 观察保留集大小从 1200 掉到多少。再想想 3.5.4 节那张表里的”2 个字”—— 你要把 top-p 拉到多少才能看到它？

3.8.2 实验 3-2　顺序陷阱

同一个 (τ, top-p) 组合，切换”先截断”和”先调温”，看保留集大小怎么变。

这个实验故意做得很小——只有两个数字在跳——因为它要传达的 就是一个反直觉的事实：两个旋钮拧的是同一个东西。

动手：把 τ 拉到 0.5，看两种顺序下保留集分别是多少。 再把 τ 拉到 2.0。哪个顺序对 τ 更敏感？为什么？

3.8.3 实验 3-3　KV 缓存的工作量

这个实验把 3.4.2 节那笔账画出来。

一个 T \times T 的方格阵，横轴是”生成第几个 token”，纵轴是”位置”。 朴素写法：每一步都要把 0 到当前位置全部涂一遍——你会看到一个三角形被反复涂。 缓存写法：预填充涂一次三角形，之后每步只涂一个新格子。

旁边一个计数器实时显示两种写法各算了多少个 token。

动手：把生成长度从 20 拉到 100。两个计数器的比值从 3 倍涨到多少？ 再想想：如果长度是 1000 呢？（3.3.2 节的公式能算出来。）

3.8.4 三个实验的观察清单

每个实验都配了”动手”提示。这一节把它们收成一份可勾选的清单—— 用于课堂演示或者自学时的自检。

实验 3-1（采样实验室）

• τ = 1.0、top-p = 1.0 时，保留集是 1200（全词表）。熵是多少？
• 把 top-p 从 1.0 拉到 0.9，保留集变成 83。分布的形状变了吗？ （提示：截断只砍尾巴，不改尖峰的高低比例——但重新归一化之后会改。）
• 把 top-p 拉到 0.5，保留集应该是 2。这两个字是什么？
• 固定 top-p = 0.9，把 τ 从 1.0 拉到 0.5。保留集变了吗？ （答案：没变，因为默认是”先截断后调温”。）
• τ = 3.0 时熵是多少？它离 \log 1200 = 7.09 还差多少？

实验 3-2（顺序陷阱）

• τ = 1.0 时，两种顺序的保留集一样吗？（应该一样）
• τ = 0.5 时，“先调温”的保留集是多少？为什么这么小？
• τ = 2.0 时呢？为什么这么大？
• 找到一个 τ 值，让两种顺序的保留集相差 10 倍以上。 是多少？

实验 3-3（KV 缓存工作量）

• 生成长度 20 时，两个计数器的比值是多少？
• 长度 50 时呢？长度 100 时呢？
• 用 3.3.2 节的公式验算：长度 N、提示词 P = 7 时的理论比值是多少？
• 网格图上，朴素写法的重复涂色集中在哪个区域？ （提示：左下角——那是”早期步骤重算早期位置”，重复次数最多。）

怎么用这份清单

上课演示时：每个实验只做前两条。剩下的留给学生课后做。 一节课的时间不够做完三个实验的全部条目。

自学时：先猜答案，再拖动滑块。 猜错了比猜对了学到更多—— 因为猜错意味着你有一个需要修正的模型。

最重要的一条：清单里每一项都指向一个具体的、可观察的差异。 如果某个实验你拖完之后”看不出变化”，那不是你操作错了， 是那个参数在当前区间内确实没影响—— 确认”这个旋钮在这里是失效的”同样是有价值的信息。

3.9 失效边界：缓存写错的六种样子

KV 缓存的实现看着简单——存下来，拼上去。但“拼上去”这一步有六种写错的方式， 而且它们症状完全不同。这一节把它们逐个跑一遍， 每种都给出真实数字和排查方法。

先立规矩：判断缓存对不对，唯一可靠的办法是比 logits，不是比生成的文本。 为什么？看完第一种 bug 就明白了。

3.9.1 判据：为什么不能比文本

我们把六种实现都在同一个提示词、同一次贪心生成上跑一遍， 对比”每一步的 logits”和”每一步的采样结果”：

实现	首个 logit 有差异的步	最大 |Δlogit|	文本分歧处数
正确实现	—	0.0000	0
pos_zero	第 1 步	12.0555	5
no_cat	第 1 步	10.8548	3
no_mask	第 0 步	2.1046	0
kv_mismatch	第 1 步	17.5341	5
cat_order	—	0.0000	0

看 no_mask 那一行：logits 最大差了 2.10——比浮点误差大六个数量级， 实现显然是错的。但生成的文本一个字都没差。

这就是”看起来能跑”的陷阱

我们用那个退化的贪心输出来比对文本。因为输出是「。，。，。，」这样的复读， argmax 的位置稳如泰山——logits 被搅乱了 2.10，但最大的那个分量 仍然还是那个分量，argmax 不变，文本就不变。

如果只测文本，这个 bug 会漏过去。 而它会一直在那里， 直到某一天换一个提示词、换一个模型，logits 的差距刚好跨过了 argmax 的边界， 输出才开始不对——而且你根本不会想到是缓存的问题， 因为”缓存版本之前一直是对的”。

正确做法：测试时直接比 logits。 run_ch3_case.py 里的 _stepwise_logits() 就是干这个的—— 它保留每一步的完整 logits 向量，而不是只留 argmax 的结果。

提示一个可以推广的测试原则

这个教训不限于 KV 缓存。

任何”数学上应该恒等”的优化，测试方法都是比中间量，不是比最终输出。

优化	恒等的东西	该比什么
KV 缓存	每步的 logits	logits
量化	每步的 logits	logits + 逐层激活
算子融合	每步的 logits	logits
分布式切分	每步的 logits	logits
混合精度	每步的 logits	logits（容差放宽到 1e-3）

共同点：只比最终输出，等于让下游的 argmax / 阈值判断 替你吸收掉了所有错误。等它吸收不掉的时候，你已经找不到原因了。

本案例里 no_mask 就是活生生的例子：logits 错 2.10，文本错 0 处。

3.9.2 pos_zero：忘了偏移位置

错误写法：缓存了 K/V，但算位置嵌入时忘了加上”已经缓存了多长”：

# 正确：位置从"已有长度"开始编号
pos = torch.arange(past, past + T)
# 错误：每次都从 0 开始
pos = torch.arange(T)          # T=1 时永远是 [0]

为什么会这么写：预填充阶段 past = 0，此时 arange(0, T) 和 arange(T) 完全一样。所以预填充测试会通过。只有进入增量阶段才会暴露。

症状：首个 logit 差异出现在第 1 步（增量第一步），最大 12.06，文本分歧 5 处。

为什么差异这么大：位置嵌入是模型用来区分”这是第几个字”的唯一信息。 位置全错等于告诉模型每个新字都是句首。第 2 章讲过， 位置信息在注意力和前馈里都会用到——全部错位，输出必然崩。

注记停下来想一想 ㉗

这个 bug 有个特点：预填充测试完全测不出来。

那么反过来问：如果一个实现只有预填充阶段（比如做文本分类、 不做生成），这个 bug 会有影响吗？

（答案：不会。预填充时 past = 0，两种写法等价。 这个 bug 只在”缓存跨越多步”的场景下存在—— 也就是自回归生成。

这提示了一个排查策略：当 bug 只在”长时间运行后”出现时， 优先怀疑”状态相关的量”——位置计数器、缓存长度、步数累加器。 它们全都符合”第一次调用时恰好正确”这个特征。）

3.9.3 no_cat：每步覆盖而不是追加

错误写法：把缓存当成”上一次的 K/V”而不是”所有历史的 K/V”：

# 正确：追加
self.k = torch.cat([self.k, k], dim=-2)
# 错误：覆盖
self.k = k

症状：首个 logit 差异在第 1 步，最大 10.85，文本分歧 3 处。

为什么：模型只能看到当前这一个 token。注意力矩阵退化成 1 \times 1， softmax 之后权重恒为 1.0，输出就是 v_t 本身—— 注意力层变成了恒等映射。

注意这和”上下文长度 = 1”不完全一样：位置嵌入仍然是正确的 （pos 还在正常增长），所以模型以为自己在写第 20 个字， 但实际能看到的只有第 20 个字。这个错位让它比”上下文 = 1” 的模型表现更奇怪。

为什么只分歧了 3 处，而 kv_mismatch 分歧了 5 处

因为贪心输出的退化性质。

「。，」这类高频组合不需要上下文就能预测——就算模型只看得到当前一个字， 它也能猜出下一个大概率是「。」或「，」。 所以退化的输出部分掩盖了这个 bug 的严重性。

第 16 个 token 才出现分歧（no_cat 的文本分歧从那里开始）， 前 15 个 token 全靠”高频组合”蒙对了。

这也是”不能只比文本”的另一个侧面：文本能对，不代表模型算对了。

3.9.4 no_mask：单步时恒真的掩码

这个 bug 最反直觉，值得单独讲。

错误写法：在缓存路径里省掉因果掩码——理由是”反正只有一步， 没有未来可以屏蔽”。

这个理由是对的。 我们用两段代码验证：

no_mask 在 T_new=1 时 max|Δlogit| = 0.000e+00
         分块 4 个时 = 3.335

一次只喂一个 token 时，加不加掩码，结果完全相同——逐位相同（0.000e+00）。

为什么恒真

缓存路径下，查询位置是 q_{\text{past}}（最后一个位置）， 键位置是 0, 1, \dots, \text{past}。

因果掩码的规则是”键位置 \le 查询位置”。

而 \text{past} 本身就是最大的键位置。所以对每一个键， “键位置 \le past” 都成立——掩码矩阵全是 1。

一个全 1 的掩码，乘不乘都一样。

那它什么时候会爆？ 当一次前向要处理多个 token 的时候：

场景	T_{new}	掩码是否恒真
单步解码	1	是（恒真，bug 不显形）
预填充	P（提示词长度）	否
分块预填充	块大小	否
投机解码验证	草稿长度	否

实测：分块大小 4 时，no_mask 的最大差异是 3.335。

而本案例的表里，no_mask 的首个差异出现在第 0 步—— 那正是预填充（一次喂 7 个 token）。所以它其实是在预填充阶段就错了， 只是后续单步全部恒真，错误没有继续扩大。

一个真实的推论

如果某个推理框架只做单步解码，而且从来不重新预填充， 那它有没有掩码都不影响结果。

但如果它支持： - 前缀缓存复用（把多个请求的公共前缀共享一份 KV）—— 拼接不同长度的前缀时，掩码必须正确 - 投机解码（一次验证多个草稿 token）—— 立刻爆 - chunked prefill（长提示词分块喂）—— 立刻爆

所以这个 bug 的危险性在于：它在你加入任何一个”提高吞吐”的特性时才会爆发。 而那时候你已经把掩码这件事忘了。

注记停下来想一想 ㉘

上面说”单步时掩码恒真”。那么：如果模型没有因果掩码（双向注意力）， 单步解码时还需要掩码吗？

（答案：需要，而且是完全不同的一种掩码。

双向模型不能做自回归生成——它需要看到整句话才能算表示。 如果硬要用它做生成（像 T5 的 encoder-decoder 那样）， 必须明确告诉它”哪些位置是真实的、哪些是 padding”。 这时掩码的作用从”防作弊”变成了”标记有效长度”，是完全不同的东西。

注意 BERT 这类模型做的是 masked language modeling——挖空填词， 不是下一个词预测。它天然不需要因果掩码。任务定义决定了要不要掩码。）

3.9.5 kv_mismatch：K 和 V 配错了对

错误写法：拼接时 K 用了倒序、V 用了正序（或者更常见的： 在一处做了 .transpose() 而另一处忘了）：

self.k = torch.cat([k, self.k], dim=-2)      # 倒序
self.v = torch.cat([self.v, v], dim=-2)      # 正序 ← 不一致

症状：这是六种里最严重的——首个 logit 差异在第 1 步， 最大 17.5341，文本分歧 5 处。

为什么最严重：注意力是 \text{softmax}(QK^\top) V。 softmax 算出的权重是”第 i 个查询对第 j 个键的关注度”。 这个权重随后会被拿去加权第 j 个值。

键和值必须一一对应。 K 的顺序变了而 V 没变， 等于说”我认为该关注第 3 个位置，但实际取到的是第 5 个位置的值”。 每一对都是错的。

为什么”K/V 都倒序”反而不是 bug

这是本节最值得记的一个发现。

我们最初把这个也列为 bug：把 K 和 V 都倒序拼接。 直觉上”历史被倒过来了”，显然是错的。

实测差异：0.0000。

一行代码的差别都没有。原因是：

\text{Attn} = \text{softmax}(QK^\top)V \tag{3.26}

注意力对键的排列是不变的。把它推导一遍。

设 P 是一个 T \times T 的置换矩阵（每行每列恰有一个 1）。 “把 K 和 V 的行按同一个置换 P 重排”就是 K' = PK，V' = PV。

（这里 K, V \in \mathbb{R}^{T \times d}，行是 token 位置，列是特征维度。 缓存拼接顺序变了，变的就是这个”行的顺序”。）

先看打分矩阵：

S' = Q K'^\top = Q (PK)^\top = Q K^\top P^\top = S P^\top \tag{3.27}

所以这个式子在说：键换了顺序，打分矩阵的列跟着换顺序。

再看 softmax。它是逐行做的——第 i 行只依赖 S 的第 i 行。 转置右乘 P^\top 只是把第 i 行的元素重新排列，不跨行混合。 由于 softmax 对行内元素的排列是等变的（分子分母同时被置换），

A' = \text{softmax}(S P^\top) = A P^\top \tag{3.28}

最后乘值矩阵：

A' V' = (A P^\top)(P V) = A \underbrace{(P^\top P)}_{=\,I} V = A V \tag{3.29}

结论：键和值按同一个置换重排，输出逐位不变。

关键在最后一步的 P^\top P = I——两次置换必须一致。 cat_order 把 K、V 都倒序，满足这个条件，所以无害； kv_mismatch 只倒序 K，P^\top 和 P 对不上，所以全错。

所以”K/V 都倒序”是一个数学恒等式，不是优化，也不是 bug。

但”只倒序 K”就不是了——因为它破坏了 K 和 V 的对应关系。 一个是恒等式，一个是全错，差别只在一个字母。

提示我们从中学到什么

我们本来准备讲”五种 bug”。数据告诉我们只有四种—— cat_order 是数学上等价的写法。

这件事本身比 bug 更有价值：它说明”看起来可疑”和”实际是错的” 是两件独立的事。判断一个实现对不对，要看数学，不能看直觉。

同样地，no_mask 在单步下差异为 0——它在这个特定配置下不是 bug， 换一个配置才是。

注记停下来想一想 ㉙

cat_order 之所以无害，是因为注意力对键的排列不变。

那么在滑动窗口注意力里（只保留最近 W 个 token 的 K/V）， 如果拼接顺序错了，还有害吗？

（答案：有害。滑动窗口本身要丢弃最旧的 K/V—— “哪个是最旧的”依赖于顺序。顺序错了，丢掉的就不是最旧的， 而是最新的——模型会失去最近的上下文。

同一个操作（换顺序），在”全注意力”下无害，在”窗口注意力”下有害。 因为窗口注意力多了一个依赖顺序的操作：丢弃。）

3.9.6 上下文越界：位置嵌入表只有那么多行

这个不是缓存的 bug，但同属”推理时才会遇到”的一类问题。

本模型的 pos_emb 是一个 (64, 128) 的嵌入表。生成超过 64 个位置会怎样？

在 CUDA 上，你会得到这个报错：

/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1500:
indexSelectSmallIndex: block: [0,0,0], thread: [64,0,0]
Assertion `srcIndex < srcSelectDimSize` failed.
CUDA error: device-side assert triggered

这个报错没有提到位置嵌入、没有提到上下文长度， 只说”某个索引超出范围”。而且它会污染整个 CUDA 上下文—— 后面的所有 CUDA 操作都会失败，必须重启进程。

怎么定位这类报错

indexSelectSmallIndex 是嵌入层查表越界的典型签名。 PyTorch 的 nn.Embedding 本质是一个 index_select。

排查顺序：

1. 哪个 embedding？ 模型里只有两个：tok_emb（词表，1200 行） 和 pos_emb（位置，64 行）。词表越界说明 token id 有问题； 位置越界说明位置索引有问题。
2. 索引是多少？ 报错里的 thread: [64,0,0] 提示了越界的索引值： 64。而 pos_emb 只有 64 行（下标 0–63）。对上了。
3. 为什么是 64？ 因为提示词 7 个字 + 生成了 57 个， 下一个位置就是 64。

注意：在 CPU 上同样的错误会给出清楚得多的报错 （IndexError: index out of range in self）。 遇到 CUDA 的 device-side assert，先用 --cpu 跑一遍， 这是最省时间的排查手段。

正确的做法是什么？ 两条路：

1. 硬限制：生成到 block_size 就停，或者提示用户。 run_ch3_case.py 里用的就是这个——在越界之前抛一个说人话的 ValueError。
2. 滑动窗口：只保留最近 W 个 token 的 K/V， 并把位置重新从 0 编号。这样上下文可以无限长， 但模型永远只看得到最近 W 个 token——这是一个能力上的妥协， 不是免费的。

警告滑动窗口为什么不是”免费的无限上下文”

有人会想：那就把窗口设大一点不就行了？

不行。 因为位置嵌入是学出来的绝对位置。 模型见过的是”位置 0 到位置 63”，如果窗口滑动之后 位置重新从 0 编号，那位置 0 在模型眼里永远是句首—— 它会以为每个窗口的开头都是新的句子。

这就是第 2 章讲 RoPE 时提到的那个问题的另一面： 绝对位置编码无法外推。RoPE 的解决方案是把位置做成相对的旋转， 窗口滑动后位置编号仍然有意义。第 4 章会展开。

3.9.7 排查手册：症状 → 嫌疑 → 验证

把这一节的内容收成一张表。排查的顺序按”排除成本”排—— 先做最便宜的检查。

症状	先怀疑	怎么验证	成本
首次生成就和朴素版不同	dim 写错 / 位置偏移	打印缓存形状，看时间维是否每次 +1	1 分钟
前 N 步相同，第 N+1 步分叉	位置计数器 / 掩码边界	打印每步的 pos，和朴素版对比	5 分钟
logits 差 10^{-6} 量级	浮点误差，正常	看是否随步数爆炸	1 分钟
logits 差 >1 但文本一样	有 bug，被 argmax 掩盖了	直接比 logits（别比文本）	5 分钟
CUDA device-side assert	位置嵌入越界	--cpu 重跑，拿清楚报错	2 分钟
输出质量莫名下降、不报错	静默截断 / 位置回绕	打印实际看到的历史长度	10 分钟
引入投机解码后接受率极低	验证路径的掩码	分块前向 vs 单步前向比 logits	15 分钟
引入前缀共享后出错	掩码不再恒真	拼接不同长度的缓存，比 logits	30 分钟

两条通用规则

规则一：先用 --cpu 跑一遍。

CUDA 的报错经常是异步的、缺上下文的、而且会污染设备。 --cpu 版本的报错通常清楚得多（比如 IndexError: index out of range in self 而不是 indexSelectSmallIndex assertion failed）。

这条规则能省掉一半的排查时间。

规则二：把”理论上不该发生”变成断言。

本章的 forward_step() 里有这么一行：

if past + T > model.cfg.block_size:
    raise ValueError(f"超出上下文长度：{past}+{T} > {model.cfg.block_size}")

如果没有它，读者会看到一段 CUDA 的错误信息，然后花半小时搜索 “indexSelectSmallIndex”——而这条搜索结果会把他引向 embedding 层的各种无关问题。

每一行断言都是在替未来的自己省时间。

一个容易走错的方向

排查数值不一致时，一个很自然的反应是：“是不是我的浮点精度不够？” 然后去改成 float64 再试一遍。

这个方向在 90% 的情况下是浪费时间。

判据在 3.4.3 节那张表里：

• 差异 < 10^{-4}：是浮点误差，不用查
• 差异 > 1：绝对不是浮点误差，查代码

本章实测的四种真 bug，最小的是 2.10——比浮点误差大六个数量级。 没有一种浮点精度问题能产生这么大的差异。

改成 float64 只会让 bug 的数值从 2.10 变成 2.1000000001， 不能解决任何问题。

3.9.8 六种实现的完整对照

把这一节的六种实现并排放。建议把这张表存下来—— 它是一份”缓存实现在不同错误下长什么样”的参考。

实现	一行代码	最大 |Δlogit|	文本分歧	危险等级
正确	—	0.0000	0	—
pos_zero	pos = arange(T)	12.06	5 处	高（预填充测不出）
no_cat	self.k = k	10.85	3 处	高（前 15 步能蒙对）
no_mask	省掉 masked_fill	2.10	0 处	极高（文本全对）
kv_mismatch	K 倒序、V 正序	17.53	5 处	高（立刻崩）
cat_order	K、V 都倒序	0.0000	0	无（不是 bug）

危险等级怎么定的

不是按差异大小，是按”能不能被发现”。

• kv_mismatch 差异最大（17.53），但它立刻让输出崩掉—— 你第一天就会发现问题。高差异 + 立刻暴露 = 容易修。
• no_mask 差异最小（2.10），但文本完全一样。 如果你只测文本，它 100% 漏过。低差异 + 不暴露 = 最难修。

所以”危险等级”= 差异大小 × 隐蔽程度，而隐蔽程度更重要。

cat_order 单独一档：它的差异是 0，因为它在数学上是恒等式。 它不是”隐蔽的 bug”，它根本不是 bug。

提示三句话总结这一节

1. 比 logits，不比文本。 六种实现里有两种在文本上完全看不出问题。
2. 形状检查能挡一半，数值比对挡另一半。 dim 写错会暴露形状， 位置和对应关系错不会。
3. 不是所有可疑的写法都是 bug。 cat_order 看着最可疑，实际最安全—— 因为 P^\top P = I（式 式 3.29）。

第 3 条最值得记住：排查时先问”这个错误在数学上会造成什么后果”， 再去改代码。本章为了确认 cat_order 无害，花的时间比修那四个真 bug 还多—— 但这个时间是值得花的，因为如果不确认，我们会在书里写一个错误的结论。

3.9.9 这一节没有覆盖的两种失效

上面六种都是缓存本身写错。还有两类失效不在缓存的代码里， 但表现出来和缓存 bug 一样，值得一并列出。

失效七：模型权重加载错。

症状：生成的文本”像是另一个模型写的”，但程序不报错。 排查：比较 model.state_dict() 的键名和检查点的键名—— 多一个 module. 前缀（DataParallel 的痕迹）、 或者某一层被随机初始化了，都会造成这个现象。

验证方法：把权重全设成 0，看输出是否变成”均匀分布的随机采样”。 如果还是像模像样的文本，说明有别的权重在起作用—— 可能是加载了两份、或者根本没加载。

失效八：tokenizer 和权重不匹配。

症状：输出是乱码，或者模型”答非所问”。 排查：把一段文本 encode 再 decode 回来，看是否一致。

assert tok.decode(tok.encode(s)) == s, "分词器不可逆"

本章的字符级分词器不满足这个断言——因为 <unk> 会丢掉信息。 这正是 3.1.2 节那个”吐标点”现象的一部分原因。

警告为什么这两类容易和缓存 bug 混淆

因为它们都是”输出不对但程序不报错”（见 3.6.3 节的静默降级）。

区分方法：缓存 bug 一定会让 logits 变化， 而权重加载错、分词器错不会——它们是”输入就错了”。

所以排查顺序应该是：

1. 先确认输入对：encode/decode 往返，token id 打印出来看
2. 再确认权重对：state_dict 键名比对，参数量对得上
3. 最后查缓存：比 logits

这个顺序不是随意定的——它按”排除成本”从低到高排。 前两步各只要一分钟，第三步要写测试代码。

本章 3.9.7 节那张排查手册的第一列就是这个顺序。

还要补一句：这三步里，第一步最容易被跳过，因为它太简单了。 但”把 token id 打印出来看一眼”这个动作， 在本章的实际排查中至少省掉过一次半小时的错误方向—— 当时我们怀疑缓存的位置偏移，实际上问题出在提示词里有一个 不在字表里的字，被 encode 成了 <unk>， 导致模型看到的前文和我们的预期完全不同。

输入没确认之前，不要开始怀疑内部实现。

3.10 上手实现：一个能跑的分步生成器

前面八节讲完了机制。这一节把散落各处的代码收成一个完整的实现， 并逐行讲解三个关键处。

完整代码在 notebooks/kv_cache.py。核心是三个东西：

1. KVCache —— 每层一份 K/V，只追加不修改
2. forward_step() —— 单步前向，可以带缓存、也可以不带
3. sample_next() —— 从 logits 到下一个 token

3.10.1 缓存本身：为什么只有 20 行

class KVCache:
    """每层一份 (k, v)，形状 (B, n_head, T_cache, d_head)。"""

    def __init__(self, n_layer):
        self.n_layer = n_layer
        self.k = [None] * n_layer        # 每层一个张量
        self.v = [None] * n_layer

    @property
    def length(self):
        return 0 if self.k[0] is None else self.k[0].shape[-2]

    def append(self, layer, k, v):
        if self.k[layer] is None:
            self.k[layer], self.v[layer] = k, v
        else:
            self.k[layer] = torch.cat([self.k[layer], k], dim=-2)
            self.v[layer] = torch.cat([self.v[layer], v], dim=-2)
        return self.k[layer], self.v[layer]

第一个关键处：dim=-2。

形状是 (B, n_head, T, d_head)。要在时间维上拼接， 也就是第 2 维（从后往前数第 2 维，即 -2）。

如果写成 dim=-1，你会在特征维上拼接—— 缓存的形状会从 (B, h, T, d_head) 变成 (B, h, 1, T*d_head)。 这个错误不会立刻报错（张量形状仍然合法）， 但 q @ k.transpose(-2, -1) 会得到一个 (B, h, 1, T*d_head) 的”注意力矩阵”—— 它在数学上毫无意义，但代码能跑。

提示怎么防住这种错误

加一个断言。 缓存每次追加之气，检查时间维确实增长了 1：

before = self.length
... 拼接 ...
assert self.length == before + 1, "缓存时间维没有增长——检查 dim"

一行断言就能把这类”形状合法但语义错误”的 bug 挡在门外。 凡是张量形状可以被多种理解的地方，都值得加断言。

3.10.2 单步前向：位置从哪来

def forward_step(model, idx, cache=None, *, bug="none", n_past=None):
    B, T = idx.shape
    past = cache.length if cache is not None else 0

    if past + T > model.cfg.block_size:
        raise ValueError(f"超出上下文长度：{past}+{T} > {model.cfg.block_size}")

    pos = torch.arange(past, past + T, device=idx.device)
    x = model.tok_emb(idx) + model.pos_emb(pos)[None, :, :]
    ...

第二个关键处：past = cache.length。

这是整个实现里最容易写错的一行，也是 3.9.2 节 pos_zero 那个 bug 的来源。

注意它在两个阶段的行为完全不同：

阶段	cache.length	pos
预填充（第一次调用）	0	[0, 1, ..., P-1] ✓
增量第 1 步	P	[P] ✓
增量第 2 步	P+1	[P+1] ✓

预填充时 past = 0，所以”用 cache.length”和”从 0 开始” 两种写法的结果一样——这就是这个 bug 能藏住的原因。

而那个 raise 也不是装饰。它是 3.9.6 节那个 CUDA 崩溃的人话版本： 与其让 PyTorch 抛一个不带上下文的 device-side assert， 不如在这里说清楚”是位置嵌入表不够长”。

3.10.3 掩码：单步时的正确处理

    if bug != "no_mask":
        qidx = torch.arange(past, past + T, device=idx.device)
        kidx = torch.arange(Tk, device=idx.device)
        allow = (kidx[None, :] <= qidx[:, None])       # (T, Tk)
        att = att.masked_fill(~allow[None, None], float("-inf"))

第三个关键处：掩码用”位置”算，不用”下标”算。

仔细看：qidx 是查询的绝对位置（从 past 开始）， kidx 是键的绝对位置（从 0 开始）。

为什么不能直接用张量下标？ 因为增量阶段两者不同步： T = 1 时，查询的张量下标是 0，但绝对位置是 past。

如果用下标算：allow = (arange(Tk)[None,:] <= arange(T)[:,None])， 当 T = 1 时右边恒为 [0]，于是只有 kidx == 0 被允许—— 模型只能看到缓存的第一个 token。这又是一个”能跑但全错”的实现。

用绝对位置算，T = 1 时 allow 全是 True（3.9.4 节证明过）， T > 1 时得到正确的下三角。一套代码覆盖两个阶段。

注记停下来想一想 ㉚

上面这个掩码，在预填充阶段（past = 0，T = P）， allow 是什么形状、什么内容？

（答案：qidx = [0..P-1]，kidx = [0..P-1]， 所以 allow[i][j] = (j <= i)——正是式 式 3.3 那个下三角矩阵。

预填充和增量用的是同一行代码，因为掩码是拿绝对位置算的。 如果拿下标算，这两个阶段就需要两套逻辑——而两套逻辑就意味着 两处可能的 bug。）

3.10.4 采样：把三种策略串起来

def sample_next(logits, *, temperature=1.0, top_k=0, top_p=1.0,
                ban=(), greedy=False, order="trunc_then_temp"):
    lg = logits.clone()
    for b in ban:
        lg[b] = -float("inf")                 # 屏蔽指定 token

    if greedy:
        return int(lg.argmax(-1))

    def truncate(z):
        if top_k > 0:                          # ① top-k
            kth = torch.topk(z, min(top_k, z.numel())).values[-1]
            z = z.masked_fill(z < kth, float("-inf"))
        if top_p < 1.0:                        # ② top-p
            srt, idx = torch.sort(z, descending=True)
            cum = F.softmax(srt, -1).cumsum(-1)
            cut = int((cum > top_p).nonzero()[0])
            m = torch.full_like(z, float("-inf"))
            m[idx[:cut + 1]] = z[idx[:cut + 1]]
            z = m
        return z

    if order == "temp_then_trunc":             # ③ 顺序
        probs = F.softmax(truncate(lg / temperature), -1)
    else:
        probs = F.softmax(truncate(lg) / temperature, -1)
    probs = probs / probs.sum()
    return int(torch.multinomial(probs, 1))

三处要点：

① 用 -inf 而不是 0。 截断掉的位置填 -inf，softmax 之后 e^{-\infty} = 0，概率严格为 0。如果填 0，那是个合法的 logit—— 它对应的概率是 e^0 / \sum，不是 0，截断根本没生效。

② top-p 的边界。 (cum > top_p).nonzero()[0] 找的是 第一个让累积概率超过阈值的位置，所以保留的是 idx[:cut+1]—— 包含那一个。差一个下标，保留集就少一个元素。

③ order 参数就是 3.5.5 节那个陷阱。 默认 trunc_then_temp （先截断后调温），因为这样 τ 和 top-p 的作用相互独立，更好预测。

警告一个必须做的最后一步

probs = probs / probs.sum()

这行不能省。

理论上 softmax 的输出已经和为 1 了。但截断之后， 被 -inf 掩盖的位置概率是 0，其余位置的概率和不再是 1 （softmax 是在截断后的 z 上做的，所以其实是对的）……

等等——truncate() 返回的 z 是截断后的 logits， F.softmax(z / τ) 自然归一化。所以这行确实是冗余的。

但它不是没用：如果 truncate 因为某种原因把所有的位置都截掉了 （比如 top_k=1 且所有 logits 都是 -inf）， softmax 会输出 nan。这时 probs.sum() 也是 nan， 除法不能救回来——但这行代码给了你一个放断点的位置。

更稳妥的写法是加断言：

assert torch.isfinite(probs).all(), "概率里有 nan——检查截断逻辑"

凡是”理论上不会发生”的分支，都值得留一个断言， 因为理论上不会发生的事情，实践中总会发生一次。

3.10.5 完整流程

把三个部分接起来：

@torch.no_grad()
def generate_cached(model, idx, n_new, *, temperature=1.0, greedy=True, bug="none"):
    model.eval()
    cache = KVCache(n_layer=len(model.blocks))

    logits, _ = forward_step(model, idx, cache=cache, bug=bug)   # 预填充
    out = idx
    nxt = sample_next(logits[:, -1, :], temperature=temperature, greedy=greedy)

    for _ in range(n_new - 1):
        logits, _ = forward_step(model, nxt, cache=cache, bug=bug)   # 每步 1 个
        out = torch.cat([out, nxt], dim=1)
        nxt = sample_next(logits[:, -1, :], temperature=temperature, greedy=greedy)

    return out

注意这里的结构对称性：预填充和增量调用的是同一个 forward_step， 区别只是传进去的 idx 长度不同（P vs 1）。

这是好设计的标志：如果一个实现需要为预填充和增量写两套代码， 那两套代码就有两个出错的地方——而这正是 3.9 节那些 bug 的温床。

提示一份可以带走的清单

写任何”数学上应该恒等”的优化时，按这个顺序做：

1. 先把正确版本写出来（朴素、慢、但显然对）
2. 写优化版本
3. 比 logits，不比输出 —— 用第 0 步到第 N 步的完整中间量
4. 设定容差 —— 浮点误差应该在 10^{-5} 以下（见 3.4.3 节）
5. 故意注入错误，验证测试能抓住 —— 这一条最容易被跳过

第 5 条是关键。如果一个测试不管你写什么代码都通过，那它没在测任何东西。 本章的 run_ch3_case.py 里那五种 bug 实现， 存在的意义就是证明这套测试真的能抓到东西。

3.10.6 一个完整的可运行例子

把前面所有零件装起来。这段代码可以直接粘进 notebook 运行。

import torch, sys
sys.path.insert(0, "notebooks")
from kv_cache import load_model, make_tokenizer, KVCache, forward_step
from tiny_gpt import Config

# ── 1. 读模型 ──────────────────────────────────────────────
DEV = "cuda:1" if torch.cuda.is_available() else "cpu"
model, itos, cfg = load_model("scaled", DEV)
tok = make_tokenizer(itos)
print(f"参数量 {model.n_params():,}　上下文 {cfg.block_size}")

# ── 2. 准备提示词 ──────────────────────────────────────────
prompt = "春江潮水连海平"
ids = torch.tensor([tok.encode(prompt)], device=DEV)

# ── 3. 采样器（3.10.4 节的版本）──────────────────────────────
@torch.no_grad()
def sample(logits, *, temperature=1.0, top_p=1.0, ban=(1,)):
    lg = logits.clone()
    for b in ban:
        lg[b] = -float("inf")                # 屏蔽 <unk>
    if top_p < 1.0:
        srt, idx = torch.sort(lg, descending=True)
        cum = torch.softmax(srt, -1).cumsum(-1)
        cut = int((cum > top_p).nonzero()[0])
        m = torch.full_like(lg, -float("inf"))
        m[idx[:cut + 1]] = lg[idx[:cut + 1]]
        lg = m
    probs = torch.softmax(lg / temperature, -1)
    return int(torch.multinomial(probs, 1))

# ── 4. 带缓存的生成 ────────────────────────────────────────
@torch.no_grad()
def generate(prompt_ids, n_new, *, seed=1, **kw):
    torch.manual_seed(seed)                  # 可复现
    cache = KVCache(len(model.blocks))
    logits, _ = forward_step(model, prompt_ids, cache=cache)
    out = prompt_ids[0].tolist()
    for _ in range(n_new):
        nxt = sample(logits[0, -1], **kw)
        out.append(nxt)
        logits, _ = forward_step(
            model, torch.tensor([[nxt]], device=DEV), cache=cache)
    return out

# ── 5. 三种解码对比 ────────────────────────────────────────
# 温度趋近 0 就是贪心（τ 不能等于 0，见 3.5.2 节的想一想）
for name, kw in [("贪心", dict(temperature=1e-6)),
                 ("τ=0.8", dict(temperature=0.8)),
                 ("τ=1.2", dict(temperature=1.2))]:
    ids_out = generate(ids, 50, **kw)
    print(f"{name:6s} {tok.decode(ids_out)}")

运行结果（本机实测）：

贪心    春江潮水海平。不見昔年年不見，不見春風吹。不見君不得，不見君家。
        不見君不見，不見，不如何處？何時？不見君不見，不
τ=0.8   春江潮水海平，狂唱山城秋。作城外由流風日裏，北樓殘月在晴雲。
        葉春無一日來，攀折綠江中。清誰與一留人，一來少年春。
τ=1.2   春江潮水海平梅水盡山春浦。君城外由還見溪裏，北陌殘江在浪飛。
        葉春無十夫江，攀折綠江柳。清誰與狂江三何處？少裁春月

注意这一段的”贪心”和 3.1.2 节那个不同：这里屏蔽了 <unk> （ban=(1,)），所以它是”念咒”版本，不是”吐标点”版本。 τ = 0.8 和 τ = 1.2 两行则和 3.1.2 节逐字相同—— 因为那两个配置在案例脚本里也没有开 top-p。

警告三个容易写错的地方

① torch.manual_seed(0) 必须放在循环外面。

放在 generate() 内部是对的——每次调用都重置种子， 所以同一个 (prompt, 参数) 组合永远给同样的输出。 这对教学很重要：读者能复现书里的每一个结果。

如果放在循环里面，每次采样都重置种子， 那么每一步抽的都是”种子 0 的第一个随机数”—— 输出会变成确定性的重复。又回到了循环。

② 屏蔽 token 用 -inf 不是 0。

lg[b] = 0 是错的——0 是一个合法的 logit， 它对应的概率是 e^0 / \sum \ne 0。屏蔽根本没生效。

③ top-p 之后必须重新 softmax。

看代码：torch.softmax(srt, -1) 是在排序后的完整 logits 上做的， 用来找截断点。而最后 torch.softmax(lg / temperature, -1) 是在截断后的 logits 上做的。

如果直接把第一次的 cum 当作最终概率，那就错了—— 因为截断改变了归一化常数。

把这段代码和第 1 章的做对照

第 1 章的案例脚本 run_ch1_case.py 里，生成是这么写的：

for _ in range(n_new):
    logits, _ = model(out)          # 每步重算整段前缀
    out = torch.cat([out, logits[:, -1, :].argmax(-1, keepdim=True)], dim=1)

第 1 章用的是朴素写法。 那是对的—— 第 1 章要讲的是 √d_k，缓存会分散注意力。

第 3 章换成缓存版本，但生成的结果逐位相同（3.4.6 节验证过）。 所以如果你把第 1 章的脚本换成缓存版本，那一章的所有结论一个字都不用改。

这是”数学恒等”这个性质的实际价值： 它让你可以在任何时候把慢的实现换成快的，而不必重跑实验。

3.10.7 真实框架怎么组织这段代码

本章的实现是教学版——清楚、可读、但性能一般。 真实的推理框架（vLLM、TensorRT-LLM、llama.cpp）在结构上和它一样， 只是把每个部分做到了极致。这一节看它们在哪里下功夫。

差异一：缓存的存储方式。

教学版用 torch.cat 每次重新分配一大块内存并把旧数据拷过去。 生成 1000 个 token 就要拷 1000 次，拷贝量是 O(N^2)—— 缓存省下的计算量，被拷贝开销吃回去了一部分。

框架的做法：预分配一块足够大的连续显存，用”写指针”追加。

# 教学版：每次都新建张量
self.k = torch.cat([self.k, k], dim=-2)

# 框架版：预分配 + 原地写入
self.k_buf[:, :, ptr:ptr+1, :] = k       # 直接写进去
ptr += 1

O(1) 的追加，没有拷贝。这也是本章 3.4.4 节”缓存没变快”的一个原因—— 教学版的 torch.cat 把省下来的时间又花掉了。

提示这个优化对本章结论有影响吗

有，但不改变结论的方向。

预分配 + 原地写入会让缓存的实测加速比更好看一些 （尤其在 CPU 上），但不会改变”GPU 上瓶颈是 kernel 启动”这个事实—— 因为启动开销和 torch.cat 是两笔独立的开销。

本章没有实现这个优化，原因有两个：

1. 教学清晰度：torch.cat 一眼就能看出”追加”这个语义， 预分配 + 指针的写法需要读者同时跟踪两块状态。
2. 它不改变任何结论：本章所有关于”缓存是恒等变换”的推导和测试 都和存储方式无关。

这是一个有意识的取舍，不是遗漏。 如果读者要拿这份代码做性能测试， 请先把它换成预分配版本——那时候数字才有意义。

差异二：注意力算子的融合。

教学版把注意力拆成 5 步：打分 → 缩放 → 掩码 → softmax → 加权求和。 每一步都是一次独立的 kernel 启动，中间结果都要写回显存再读出来。

框架用 FlashAttention 把它们融成一个 kernel： 中间结果留在片上（SRAM），不落显存。 这一项能给长序列带来 2–4 倍提速，而且省显存。

差异三：连续批处理（continuous batching）。

教学版一次处理一条请求。框架把不同请求的”第 k 步”动态拼成一个 batch：

时刻 t：  [请求A的第5步] [请求B的第2步] [请求C的第9步]
时刻 t+1：[请求A的第6步] [请求B的第3步] [请求D的第1步]   ← B 结束了，D 进来

关键是”动态”：不需要等整个 batch 都结束才换人。 这解决了传统 batching 的”长尾问题”—— 一条请求生成了 5000 个 token，其他请求早结束了， 传统 batch 要一直空等。连续批处理把 GPU 利用率从 30% 提到 90% 以上。

三项优化分别对应哪一层

优化	解决什么	对应本章哪一节
预分配 + 写指针	拷贝开销	3.4.5 显存管理
FlashAttention	kernel 启动 + 显存带宽	3.3.3 compute/memory-bound
连续批处理	权重复用率	3.3.3 “batch 是唯一的解药”

注意三项都不是新算法。 它们都是工程—— 但每一项都建立在对”瓶颈到底在哪”的准确判断上。

本章 3.4.4 节那个”省 68 倍计算只快 1.07 倍”的实测， 价值就在于它告诉你瓶颈不在计算。 如果不知道这一点， 你可能花几个月去优化矩阵乘法，而真正的瓶颈在 kernel 启动。

警告一个反过来的教训

上面说框架的实现在每一层都比教学版强。 那么直接抄框架的代码来学习，行不行？

不行，因为框架的代码是”结论”，不是”推理”。

看 vLLM 的 attention.py，你会看到分页、写指针、CUDA 图、 自定义 kernel 交织在一起。你无法从这段代码看出”为什么要缓存”—— 因为”为什么”已经在几百次 commit 之前被决定了。

教学代码和产品代码的目标不同：

	教学版	产品版
目标	让人看懂	跑得最快
长度	20 行	2000 行
性能	慢	快
可读性	高	低

先用教学版建立心智模型，再去看产品版对照”它优化了什么”。 反过来做，你会迷失在细节里而不知道每个细节在解决什么问题。

3.11 延伸阅读

本节五篇，按”读它能得到什么”排序，不是按重要性。

3.11.1 The Curious Case of Neural Text Degeneration (2020)

Holtzman, Buys, Du, Forbes, Choi. ICLR 2020。 arXiv:1904.09751。

读它能得到什么：本章 3.1 节那个”贪心解码必然循环”的现象， 这篇论文是系统研究它的第一篇。作者在 GPT-2 上做了本章做过的实验 （贪心 / beam search / 纯采样 / top-k / top-p 的对比）， 并提出了 nucleus sampling（top-p）。

读的时候注意什么：

• 第 3 节是核心。作者指出 beam search 的失败不是因为 “搜索不够充分”，而是因为最大化似然本身就是错的目标。 这个论断和本章 3.3.5 节的分析是同一件事的两种说法。
• 图 1 值得盯着看。它画出”人类文本的概率分布”和 “机器解码文本的概率分布”——前者有起伏，后者单调下降。 这个对比比任何文字描述都有说服力。
• 注意它的局限：这篇论文的实验全是英文， 而且用的是 GPT-2 级别的模型。本章用一个 42 万参数的 中文模型复现了同样的现象——说明这个现象和规模无关。

3.11.2 The Curious Case of Neural Text Degeneration 之前的铺垫：Mirostat

Basu, Ramachandran, Keskar, Varshney (2021). Mirostat: A Neural Text Decoding Algorithm that Directly Controls Perplexity. ICLR 2021。arXiv:2007.14966。

读它能得到什么：top-k 和 top-p 都是间接控制——你调的是 k 或 p， 但你真正想要的是”输出有多出人意料”。Mirostat 直接把这个目标写出来： 维持一个目标困惑度，用反馈控制动态调整截断阈值。

读的时候注意什么：这篇的价值在于把”我们到底在控制什么”这个问题摆到台面上。 读完再回头看本章 3.5 节那张”顺序陷阱”的表， 你会意识到那不只是个实现细节——它暴露了这些旋钮的语义本身就不清晰。

3.11.3 Fast Transformer Decoding: One Write-Head is All You Need (2019)

Shazeer. arXiv:1911.02150。

读它能得到什么：本章 3.4.5 节算了 KV 缓存的显存开销 （Llama-3-70B 单序列 21.5 GB）。这篇论文提出 MQA （multi-query attention）：让所有注意力头共享同一份 K/V， 把缓存缩小到 1/h。

读的时候注意什么：

• 这是一篇只有 4 页的论文，属于”想法比篇幅重要”的那类。
• 代价在哪：MQA 会损失一些质量。后来的 GQA（grouped-query attention，Llama-2/3 用的就是它） 是 MQA 和 MHA 的折中——分组共享，组内不共享。
• 和本章的关系：3.9 节那六种缓存 bug 里， 没有一种是因为”头之间共享了 K/V”引起的。MQA 不改缓存的语义， 只改它的形状——所以第 3.9.1 节那套”比 logits”的测试方法 换成 MQA 之后仍然适用。

3.11.4 Efficiently Scaling Transformer Inference (2022)

Pope et al. arXiv:2211.05102。

读它能得到什么：这篇是推理优化的系统性综述， 把”显存带宽 / 算力 / 延迟”三个瓶颈分开讨论。 本章 3.4.4 节那个反直觉的实测结果 （省 68 倍计算，GPU 上只快 1.07 倍）在这篇里有一个完整的理论框架： 推理分两个阶段，prefill 是 compute-bound，decode 是 memory-bound。

读的时候注意什么：

• 第 2 节的两张表是全篇精华。它把不同模型规模、 不同序列长度下的瓶颈在哪里列得很清楚。
• 读完再回头看 3.4.4 节，你会明白为什么本章测到的是 1.07 倍： decode 阶段在 batch=1 时，GPU 的算力利用率只有几个百分点， 少算 68 倍根本不影响墙钟。
• 它解释了”为什么 batch 大了缓存才有用”：batch 一上去， 瓶颈从”启动开销”变成”带宽”，缓存省下的那些读写就变成实打实的时间。

3.11.5 vLLM / PagedAttention (2023)

Kwon et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023。arXiv:2309.06180。

读它能得到什么：3.4.5 节的”停下来想一想 ⑦”提到， KV 缓存的显存是动态增长的，而且长度不可预测。 PagedAttention 借用操作系统的虚拟内存分页思路解决这个问题： 把缓存切成固定大小的块（block），用一张页表把逻辑位置映射到物理块。

读的时候注意什么：

• 读第 3 节就够了（PagedAttention 的设计）， 第 4、5 节是实现细节和评测。
• 注意它解决的三个问题：① 预分配浪费（你不知道要生成多长）； ② 外部碎片；③ 前缀共享——同一个提示词的多个输出可以共享同一份前缀 KV。
• 第 ③ 点和本章直接相关：前缀共享意味着缓存的 K/V 可能来自不同的序列。这时 3.9.4 节那个”单步时掩码恒真”的结论 不再成立——因为拼接的缓存长度不等于查询位置。 这正是那个 bug 会突然爆发的地方。

---

3.11.6 Attention Is All You Need 里的推理（2017）

Vaswani et al. NeurIPS 2017。arXiv:1706.03762。第 3.2.3 节脚注所引。

读它能得到什么：本章讲的”训练并行、推理串行”这个二分， 在 Transformer 原论文里就已经明确写出来了。原文第 3.1 节：

“While the model is auto-regressive at inference time, the training can be parallelized…”

读的时候注意什么：

• 注意原文说的是”推理时自回归”，而不是”推理时慢”。 2017 年的论文没有讨论 KV 缓存——那篇论文的贡献是架构， 推理优化是后来几年的工程工作。
• BERT（2018）之后，很多人误以为”Transformer 不是自回归的”。 实际上编码器和解码器是两种不同的东西： BERT 是编码器（双向、非自回归、不能生成）， GPT 是解码器（因果、自回归、能生成）。 本章讲的所有内容都只适用于解码器。

3.11.7 The Illusion of Consensus（2024）——一个警示

Atil et al. arXiv:2405.00662。

读它能得到什么：这篇论文做的是复现性研究—— 把一批”best-of-n 采样能提升大模型推理能力”的论文拿来重跑， 发现大部分结论在控制住”评分方式”之后消失了。

读的时候注意什么：

• 它和本章的直接关系：本章 3.5.9 节说”数学推理常用 best-of-n”。 这篇论文正好质疑这个做法——如果评分器（reward model） 有偏好，best-of-n 只是在放大那个偏好，不是在找更好的答案。
• 它是一个方法论警示，不是结论。 读完不应该得出 “best-of-n 没用”，而应该得出”声称 best-of-n 有用时， 必须说清楚评分器是什么、怎么排除评分器的偏好”。
• 这一条对写论文和读论文都适用。 本书反复强调的 “数字必须真跑出来、结论必须带任务前提”， 在这篇论文里被推到了极致——连别人跑出来的数字都要重跑一遍。

---

这一章真正要你记住的是

“训练”和”推理”不是同一个算法的快慢两个版本，它们是两件事。

训练时，模型在真实前文下被要求预测下一个字， 一次算完所有位置，用交叉熵衡量。这是它的能力。

推理时，模型在自己造的前文下逐字生成， 一次只能算一个位置，用解码策略把分布变成一个具体的字。这是它的行为。

能力到行为之间，隔着一个解码策略。 本章所有的现象——循环、退化、顺序陷阱、置信度错位—— 都不是能力问题，是这一层的问题。

而 KV 缓存是唯一一个跨越两副面孔的东西： 它让推理用上训练的并行性，代价是必须证明它恒等—— 因为一旦不恒等，它产生的就是一个你无法用 loss 衡量的错误。