第 4 章 推理生成¶
第 3 章的 generate 已经能一个 token 一个 token 往外写，但每一步都把当前前缀整段再送进网络重算一遍。要把延迟和并发做上去，需要分清 Prefill 与 Decode，并把历史的 Key、Value 缓存起来。
4.1 Prefill 与 Decode¶
一次在线生成可以分成两段。
Prefill。 输入是已经写好的上下文：系统提示、历史对话、当前用户消息等，长度 \(n\) 已知。模型对这 \(n\) 个位置做一次前向（因果掩码下各位置可并行），得到第一个要回复的 token。用户感知的「首字延迟」，主要落在这一段。
Decode。 从第二个生成 token 起，每步只增加一个新位置：把该位置的 id 送进网络，跑完 \(L\) 层，取样得到下一个 id，拼回序列。下一步的输入依赖这一步的输出，所以在序列长度方向上必须串行。逐字延迟由这一段决定。
第 3 章的循环把两段混在同一种调用里：每一步都对「当前整段前缀」做完整 forward。接上 KV 缓存之后，Decode 才变成「只算新位置」。
4.2 KV 缓存¶
为何要缓存¶
Decode 时，新 token 在每一层注意力里仍然要和左侧全部历史位置做点积。历史位置的 Key、Value 在更早的步里已经算过；在因果模型里，过去位置的表示不会因为右边新写出一个 token 而改变。若像第 3 章那样每步从第一个位置重新算到当前末尾，同一段历史的 \(K\)、\(V\) 会被算了又算，计算量大致随已生成长度累加。
因此把各层历史位置的 \(K\)、\(V\) 留在显存里，称为 KV 缓存。Decode 一步只对新 token 算 \(Q,K,V\)：新的 \(K,V\) 追加进缓存，新的 \(Q\) 去和缓存里全部历史 \(K\) 算注意力，再加权历史 \(V\)。当前步的 \(Q\) 用完即弃，不必写入缓存——后续步不会再查询「上一步的 Query」。
Prefill 结束时，把 prompt 各位置的 \(K,V\) 写入缓存，作为 Decode 的起点。之后每生成一个 token，缓存长度 \(n\) 加一。
形状与显存¶
逻辑上可以把全模型的缓存看成：
| 维 | 含义 |
|---|---|
| \(B\) | 并发请求数（batch） |
| \(L\) | 层数；每层各有一份 K、V |
| \(n_{\mathrm{kv}}\) | 每层写入缓存的 KV 头数 |
| \(n\) | 当前上下文长度（prompt + 已生成），逐步递增 |
| \(d_h\) | 每个头的维度 |
K 与 V 各一份，元素总数为 \(2 B L n_{\mathrm{kv}} n d_h\)。设每个元素占 \(b_{\mathrm{kv}}\) 字节（bf16 下为 2），则
中间层的激活在该步算完后可以释放；KV 要保留到这条请求结束。因此 \(M_{\text{KV}}\) 随 \(n\) 与 \(B\) 线性增长。模型权重大小不随对话变长，长上下文或高并发时，显存往往先被 KV 顶满。
数量级：8B 级模型 bf16 权重大约 16 GB。第 2–3 章的多头设定里，每个 Query 头各有一份 \(K,V\)，故 \(n_{\mathrm{kv}}=h\)。取 \(L\approx 36\)、\(h=n_{\mathrm{kv}}=32\)、\(d_h=128\) 时，全模型每增加一个 token 的 KV 大约四百多 KB；\(n\) 到数万时，单请求的缓存就可以到数 GB。更细的表见附录 E。
公式里 \(M_{\text{KV}}\) 与 \(n_{\mathrm{kv}}\) 成正比，因此不必坚持 \(n_{\mathrm{kv}}=h\)。若若干个 Query 头共用同一组 \(K,V\)，缓存只存 \(n_{\mathrm{kv}}<h\) 份，显存按同一比例下降。怎样共用、质量如何，见第 7 章；读本章公式时，把 \(n_{\mathrm{kv}}\) 当成可以单独选取的因子即可。
和权重、计算量怎么一起看¶
推理显存里，和 KV 并列的是权重。记 \(P\) 为参数个数（第 3 章对 parameters() 求和的那个数；「8B」即 \(P\approx 8\times 10^9\)），每参数 \(b_w\) 字节，则
bf16 下 \(b_w=2\)，8B 权重约 16 GB。它不随 \(n\) 变长；变长的是 \(M_{\text{KV}}\)。
Decode 每步的计算量可粗估为
其中 \(d\) 为隐藏宽度（n_embd）。前一项是当前 token 走完（稠密情形下）全部参数的矩阵乘；后一项是新 Query 与长度 \(n\) 的历史做注意力。\(n\) 不大时前一项主导，所以 Decode 往往卡在搬权重；\(n\) 很大时后一项和读 KV 都会明显起来。
这里的 \(2P\) 有一个前提：每个新 token 都经过每一层的全部注意力与 FFN 参数——这种激活方式叫稠密（dense）。第 2–3 章的模型就是稠密的。若每 token 只激活一部分专家，\(2P\) 里的 \(P\) 应改成激活参数量，见第 9 章。
Prefill 一次处理长度 \(n\) 的上下文时，粗估为 \(C_{\text{prefill}}\approx 2nP+4Ln^2d\)（整段 \(n\times n\) 注意力）。附录 E 可按配置代入。
4.3 Prefill 与 Decode 各慢在何处¶
GPU 上要同时看：矩阵乘单元能不能持续算（算力），以及数据能不能及时从显存送到算力单元（带宽）。大块参数放在容量大、较慢的 HBM 里。算力单元一直有活干，耗时跟 FLOPs 走，叫算力受限；单元经常空等下一块数据，耗时跟搬移字节数走，叫访存受限。
Prefill。 权重读入后会被 \(n\) 个位置反复用，投影和 FFN 常常接近算力受限。注意力若把 \(n\times n\) 分数表整表写入 HBM 再读回，\(n\) 较大时带宽先满，Prefill 变成访存受限，峰值显存也被抬高。FlashAttention 等实现保持相同数学结果，在片上 SRAM 分块做完注意力，避免把整张分数表写入 HBM 再读回；配对仍是平方级，只是少搬那张大表。
Decode。 每步只有一个新 token，算术相对整份权重很小，时间主要花在从 HBM 搬权重，以及读长度为 \(n\) 的 KV。这是访存受限。KV 缓存消除的是历史的重复计算，并不消除「每步读权重」和「读变长的 KV」。
两段慢点不同：Prefill 要防 \(n\times n\) 的大块读写；Decode 要控权重搬运与 \(M_{\text{KV}}\)。
动手¶
固定一个稠密模型，扫描上下文长度与 batch，记录 KV 占用与 Decode 延迟，与 \(M_{\text{KV}}\)、\(C_{\text{decode/token}}\) 对照。
核心问答¶
Q1：训练时为什么不维护 KV 缓存？ 训练时整段 token 已知，一次前向里各位置的 \(K\)、\(V\) 一并算出。缓存解决的是生成时右侧尚未写出、又不能每步重算全部历史的问题。
Q2：为什么只缓存 \(K\)、\(V\)，不缓存 \(Q\)？ 后续步要用历史 Key 做匹配、用历史 Value 做加权。当前步的 Query 只用于当前步。
Q3：有了 KV 缓存，Decode 还会随长度变慢吗？ 会。每步仍要读权重，还要读长度为 \(n\) 的 KV。缓存消除的是历史的重复计算，不是长度本身。
Q4：稠密和「参数量大」是一回事吗？ 稠密指每 token 激活层内全部相关权重。总参数大但每 token 只走一部分专家时，Decode 应按激活参数估 \(2P\)。
Q5：FlashAttention 是否把复杂度降到了 \(O(n)\)？ 没有。它减少 \(n\times n\) 的 HBM 读写与峰值显存；位置配对仍是平方级。
Q6：温度、量化、投机解码、长思维链，各自改变哪一项？ 温度（及 top-\(p\)）只改 logits→token 的取样，不动 Prefill/Decode 计算图。量化压 \(b_w\) / \(b_{\mathrm{kv}}\)，减权重或 KV 的字节。投机解码把一次权重读取摊到多个待验证位置。长思维链加的是 Decode 步数与 \(n\)，因而放大 \(M_{\text{KV}}\) 与访存时间。