Step by Step 读懂 LLM（三）：分类器、Prefill、Decode 与 KV cache

这是 Step by Step 读懂 LLM 的最后一篇。前两篇我们把一段文字一路送进了模型：第一篇讲 BPE 和 embedding，第二篇讲 Transformer block 里的 LayerNorm、Self-Attention 和 MLP。这一篇讲剩下的三件事：最后怎么从一个向量变成下一个词，推理时的两个阶段 Prefill 和 Decode，以及让 Decode 快起来的 KV cache。

右边上面是一个只有两层的 LLM，每一列是一个 token，从下往上是它经过的每一个计算单元，q、k 在打分前都先做了 RoPE。prompt 是 The tired cat，模型刚生成了 crossed。读到不同的段落，图会切到对应的情况；也可以自己切换，或者点任何一个格子看它在算什么。

1分类器：从向量到下一个词

现在一个词进入 Transformer，经过了 embedding、n 层 Transformer block，最后输出了这个词在这个位置上的 feature vector，再过最后一个 RMSNorm，记作 （1×d）。接下来一步就是需要在这个位置生成一个词。

这个逻辑很简单，就是一个多分类的 logistic regression（也叫 softmax regression）。我们把 乘以一个 d×V 的矩阵 （V 是词表大小），得到词表里每个词一个分数，叫 logits；再通过 softmax 归一化，对词表内所有词生成一个 probability；然后基于这个 probability 进行随机抽样，获得这个位置的词：

乘 其实也是一次打分：它的第 i 列 可以看成第 i 个词的“输出向量”，，就是 和每个词有多像。这和 attention 里 打分是同一个套路，只不过 attention 是在前文的 token 里挑，这里是在整个词表里挑。有些模型（比如 GPT-2）直接让 ，用的就是第一篇里的 embedding 矩阵，叫 weight tying。

这一步并不便宜： 在 Llama 2 7B 里是 4096 × 32000 ≈ 1.3 亿次乘加，和一层 MLP 差不多；Llama 3 的词表有 128,256 个词，光这一个矩阵乘法就超过一整层 Transformer。好在推理时只需要最后一个位置的预测，前面位置的分类器都不用算（右图上图上方写着“不需要”的那几列）。

Cross Entropy：训练时每个位置都在考试

推理时只看最后一个位置，训练时就不一样了。训练数据里，每个位置的“正确答案”都是已知的：就是它后面那个 token。所以训练时每一列都会算出一个概率分布，然后看它给正确答案的概率有多高（右图上图切到“训练”：每一列的预测都用上了，上面标的是它要猜的那个词）。

损失函数就是 Cross Entropy（交叉熵）：正确答案的概率取负对数，再对所有位置取平均：

模型给正确答案的概率是 1，这一项就是 0；概率越小，损失越大，而且是对数式地变大，猜得离谱会被重罚。这就是这一节标题“Cross Entropy 分类器”的来历：推理时只用到 softmax，Cross Entropy 是训练时用来打分的。

还记得第二篇的 causal mask 吗？训练时整句一次并行算完，每个位置都要猜它后面的词，如果第 i 个位置能看到第 i+1 个 token，就等于直接把答案给它了。causal mask 让所有位置可以同时考试，又不会互相抄答案。

抽样：temperature、top-k 和 top-p

有了概率之后还要真正选出一个词。这里有几个参数会影响抽样，包括 temperature、top-p 和 top-k。下面每一段都可以在 图 2 ↗ 里拖动看效果，点右边上面模型最上面的 p 格子也能打开。

temperature 可以用于平滑或者激化 softmax 的曲线：softmax 之前的 logits 会先除以 temperature T，然后再走过 softmax 这一层。T 小于 1，会让高概率的词概率更高、更突出；T 大于 1，会让所有的概率更平滑。这是 exponential function 的特性：除以 T 不改变 logits 的大小顺序，只是把它们之间的差距放大或缩小，而 softmax 只看差距。T → 0 就是 greedy，永远选概率最大的那个词。

top-k 和 top-p 则是选择从哪些词里抽样。top-k 是只从概率最高的 k 个词里抽取；top-p^[2] 是把概率从高到低累加，取累加到 ≥ p 为止的最少那几个词，只从这里面采样。这样可以避免抽到一些概率极小的词，影响整个句子。两者的区别在于：top-k 永远留 k 个；top-p 留几个是自适应的，模型很确定时可能只留 1 个，分布很平时会留很多个。

常见实现的顺序是：logits ÷ T → top-k → top-p → 重新归一化 → 按概率抽一个。top-k 和 top-p 都设的时候，先 top-k、再在剩下的词里做 top-p，最后留下的是两者的交集，谁更严格谁说了算。点开图 2 ↗，把这几步一步步展开。

2Prefill 和 Decode

推理分成两个阶段。

Prefill：把整段 prompt 一次性送进模型。prompt 里的所有 token 都是已知的，所以可以像训练时一样，把 t 个 token 叠成一个 t×d 的矩阵，每一层都是一次大矩阵乘法，所有位置并行算完（右图上图切到 Prefill：三列一起算）。Prefill 做两件事：算出最后一个位置的预测，生成第一个新 token；同时把每个 token 在每一层的 k、v 存进 KV cache（绿框）。用户等多久才能看到第一个字（TTFT，time to first token），主要由 prefill 决定。

Decode：之后的 token 只能一个一个生成，因为第 t+1 个 token 要等第 t 个被抽出来，才知道输入是什么。每一步只把这一个新 token 送进模型（右图上图黄色那一列），用它的 q 去读 cache 里所有前文的 k、v，输出下一个词的概率，抽样，再把抽到的词作为下一步的输入。之后每个字隔多久出来（TPOT，time per output token），由 decode 决定。

最大的区别：一个卡算力，一个卡带宽

两者最大的区别是瓶颈不同：prefill 卡在算力上（compute-bound），decode 卡在显存带宽上（memory-bound）。

每做一次 ，GPU 都要先把权重 从显存（HBM）搬到计算单元。关键看搬一次权重能做多少计算。Prefill 是 (t×d)(d×d)：权重搬一次，t 行一起用，每个字节大约能做 t 次运算。Decode 是 (1×d)(d×d)：权重搬一次只用 1 行，每个字节大约只能做 1 次运算。这个“每字节能做几次运算”叫算术强度（arithmetic intensity）。

以 H100 为例^[3]，算力约 989 TFLOPS，带宽约 3.35 TB/s，两者相除，分界线大约是每字节 295 次运算：高于它，算力先用满；低于它，搬数据先到上限，计算单元在等数据。下图拖动一次处理的 token 数，看工作点在 roofline 上怎么移动。

Decode 每生成一个 token，都要把 13.5 GB 的权重从显存读一遍，大约 4 ms；而真正的计算只要 0.01 ms，GPU 的算力用了不到 1%。上下文越长，还要多读越来越大的 KV cache。

	Prefill	Decode
怎么算	整段 prompt 并行，一次过完	一个接一个，下一个要等上一个抽出来
做了什么	建好 KV cache，输出第一个 token	往 cache 里追加一行，每步输出一个 token
瓶颈	算力（compute-bound）	显存带宽（memory-bound）
决定哪个延迟	TTFT：等多久看到第一个字	TPOT：之后每个字隔多久

工程上怎么办

因为瓶颈不同，优化方向也不同。提升 decode 效率的思路，都是读一次权重，多干点活：

• Batching：把很多用户的 decode 放到同一批里一起算。权重读一次，B 个 token 一起用，算术强度变成约 B（图 3 里的 batch 64）。这是推理服务成本能降下来的关键。
• 投机解码：小模型先猜几个 token，大模型一次验完，相当于把 decode 变成一次很小的 prefill^[4]。我在另一篇里专门讲过。
• 量化：权重用更少的位数存，同样的带宽每次读得更快。

两者混在一起也有问题：一个长 prompt 的 prefill 会占满算力，把同一批里其他人的 decode 卡住。所以有了 chunked prefill^[5]（把长 prompt 切成小块，和 decode 穿插着算），以及把 prefill 和 decode 放在不同机器上、各自按自己的瓶颈配置资源的分离部署^[6]。

3KV cache

在一个对话或者 agent 使用的过程中，相同的上下文会在短时间内被反复拿来计算。先看一次生成的内部：prefill 之后开始 decode，如果我们去看 decode 的过程，它是在反复使用相同的前文的 k 和 v，因为每个新 token 都是用它自己的 q，去和所有前文的 k、v 做计算。这就是 KV cache 的意义：这些 k、v 的值被存在 GPU 显存里，然后被重复利用，节省计算。

如果不存会怎样？右图上图切到“Decode · 无 cache”：为了拿到前面 3 个 token 在每一层的 k、v，只能把它们从查表开始整列重新算一遍（红色），结果和上一步一模一样；它们在最后一层的 attention、MLP 和分类器更是算了也用不上（红色斜线）。一共要算 144 个单元，其中 52 个矩阵乘法。

有了 cache，只算新 token 这一列（黄色），旧 token 在每一层的 k、v 直接从 cache 读（绿色）：36 个单元，13 个矩阵乘法。而且不管前面有多少 token，都只算 1 列。换成真实长度，prompt 1000 个 token、再生成 1000 个：没有 cache 要把 token 过模型约 150 万次，有 cache 只要 2000 次。

为什么旧的 k、v 可以一直复用？这正好是第二篇里的几件事：causal mask 让每个 token 只看前面，新 token 加进来，旧 token 看不到它；LayerNorm 只在每个 token 内部做，不跨 token；RoPE 按每个 token 自己的位置旋转 k，位置不变，转好的 k 也不变。所以 cache 里存的是转过之后的 k（右图上图的 R(k)）和 v。

那为什么不存 q？旧 token 的 q 只用来算它自己的输出，早就算完了；新 token 只需要它自己的 q，加上所有人的 k 和 v。

有了 cache 之后，每一步里只有 attention 的 q·k 和 Σ α·v（加上中间的 softmax）还会随长度变长：要和 cache 里的 t 个 k、v 各算一次。对 7B 模型，大约到 t ≈ 25,000 时，这部分的计算量才追上和权重相乘的部分。真正更要紧的，是每一步都要把整个 cache 从显存读一遍：同样大约在 26,000 个 token 时，读 cache 的量就超过读权重的 13.5 GB。

代价：用显存换计算

cache 要一直放在显存里。7B 模型（32 层，d = 4096，fp16）每个 token 要存：

所以围绕 KV cache 有一整套办法：GQA / MQA 让几个头共用一组 k、v^[7]；PagedAttention 像操作系统管内存页一样管 cache，减少碎片^[8]；把 cache 量化成更少的位数；或者干脆只保留最近一段窗口。

回到开头说的“同一段上下文被反复使用”：KV cache 本身是一次生成内部的复用。到了多轮对话和 agent，复用还能更进一步：system prompt、工具说明、前面几轮对话这些前缀，在下一次请求里一字不变，它们的 k、v 也一模一样。把这部分 cache 留在显存里给下一次请求接着用，下一次的 prefill 就只需要算新加的那一小段。这叫 prefix caching，很多 API 里的 prompt caching 就是这个思路。

4把三篇串起来

到这里，一个 token 从文字到下一个词的完整路径就走完了：

1. BPE：文本切成 token，每个 token 一个 id（第一篇）。
2. Embedding：用 id 查表，得到 t×d 的矩阵（第一篇）。
3. Transformer block × L 层：RMSNorm → Self-Attention（q、k 做 RoPE，causal mask，多头）→ 残差 → RMSNorm → MLP → 残差（第二篇）。
4. 分类器：最后一个 RMSNorm，乘 W_vocab 得到 logits，softmax 得到概率；训练时用 Cross Entropy 打分，推理时按 temperature / top-k / top-p 抽样（本篇第 1 节）。
5. 推理：prompt 一次 prefill，之后一个一个 decode；KV cache 让 decode 每步只算一列，代价是显存（本篇第 2、3 节）。

每一步里，token 之间交换信息的只有 attention 一处，其他全部是每个 token 各算各的。也正因为这样，prefill 能把整句并行，decode 能只算一列，KV cache 才成立。

参考

1. Vaswani et al. Attention Is All You Need. NeurIPS 2017.
2. Holtzman, Buys, Du, Forbes, Choi. The Curious Case of Neural Text Degeneration（top-p / nucleus sampling）. ICLR 2020.
3. NVIDIA. H100 Tensor Core GPU（SXM：BF16 dense ≈ 989 TFLOPS，HBM3 3.35 TB/s）.
4. Leviathan, Kalman, Matias. Fast Inference from Transformers via Speculative Decoding. ICML 2023.
5. Agrawal et al. SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills. 2023.
6. Zhong et al. DistServe: Disaggregating Prefill and Decoding for Goodput-optimized LLM Serving. OSDI 2024.
7. Ainslie et al. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023.
8. Kwon et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023.
9. Pope et al. Efficiently Scaling Transformer Inference. MLSys 2023.