Speculative Decoding 投机解码

1Prefill 与 Decode

LLM Inference 主要分为两个阶段：Prefill 和 Decode。

Prefill 阶段，LLM 的输入是完整的 prompt，也就是我们提供的 input。模型要为每个 token 算出每一层对应的 K 和 V（也就是 KV cache），供后面 decode 阶段生成新 token 时使用。因为所有 token 都是已知的，这个过程可以并行：模型权重从显存读一遍，就能同时给几百上千个 token 做计算。计算量很大，花在计算上的时间比读显存还多，所以 Prefill 通常是 Compute Bound 的，也比较容易通过 TP、EP、DP 等各种并行方式用更多的显卡来加速。

Decode 阶段就不一样了。每次都需要用前 x 个 token 预测第 x+1 个 token；要预测第 x+2 个，就得等第 x+1 个完全算完才能开始。而生成一个 token，要把模型的全部参数从显存读进计算单元一遍。一个 700 亿参数的模型用 BF16 存储，就是 140 GB。H100 的显存带宽约 3.35 TB/s，光是搬运就要 ≈ 42 ms。

真正的计算部分呢？每个 token 大约需要 2N 次浮点运算，也就是 1400 亿次。H100 的 BF16 算力约 989 TFLOPS，只要 ≈ 0.14 ms，远小于读显存的时间。所以 Decode 通常是 Memory Bound 的。再加上逐个生成带来的线性依赖，它也没法像 Prefill 那样靠并行来加速。

一次前向的耗时roofline

一次前向处理的 token 数k 1

这一次前向的时间线

图 1拖动滑块：decode 每次只处理 1 个 token，落在曲线平坦的左端；prefill 一次处理整个 prompt，落在右边。理想化估算，假设计算与搬运完全重叠，忽略 KV cache、注意力和通信；拐点 ≈ 989 TFLOPS ÷ 3.35 TB/s ≈ 295，与模型大小无关。

"读显存"读的是哪一段？

模型权重和 KV cache 在整个生成过程中一直待在 GPU 的显存（HBM）里，并不会每一步都从 CPU 重新拷一遍。这里说的"读"，是 GPU 内部从 HBM 读进计算单元。计算单元旁边的片上存储（SRAM）全卡加起来只有几十 MB，装不下动辄上百 GB 的权重，所以每生成一个 token，全部权重都要从 HBM 流过计算单元一遍：算完一块，丢掉，再读下一块。

权重从 HBM 流进计算单元

图 2示意图，比例不真实。每个蓝色小块是一块权重，黄色火花是用它做的一次乘加。decode 时每块权重读进来只用一次；prefill 时同一块权重被很多 token 共用。

一次读了这么多数据，却只生成一个 token，
实在太不划算了。

Speculative Decoding 就是为了解决这个问题。

2从 CPU 分支预测说起

投机解码的核心想法，受到了 CPU 分支预测（branch prediction）的启发。

一个普通的 CPU 程序里有很多条件判断：if、while、switch 等等。想象一个要运行 1 万次的 while 循环，每次都要判断一次括号里的条件。正常来说，CPU 需要把条件完整算完，才能决定接下来执行哪部分代码，这会严重拖慢整个程序。

所以 CPU 做了一个优化：猜。在条件算完之前，先猜一个最可能的分支，提前开始执行；等条件算完再验证。猜对了，就白赚了这段时间；猜错了，就把提前算的结果丢掉，回到正确的分支重新执行。这里包含两部分：

一个比执行条件语句廉价得多的预测算法；
验证与修正。

分支预测的时间线

图 3猜对时，执行和条件计算重叠，整体提前结束；猜错时只损失一点回滚的开销，结果仍然正确。

回到 LLM inference。给定前 x 个 token，要预测第 x+1 到第 x+n 个 token 时，第 x+i 个 token 必须等前面 i−1 个一个一个算完才能开始，就像 while 循环的每一步都要等括号里的条件算完，才能进入下一步。

所以我们也需要一个廉价得多的"猜测算法"，再由大模型用相对昂贵的计算来验证。这就是 Speculative Decoding 在做的事情。

3算法详解

猜测。在 Google 最早的论文^[1]中，他们用同家族但小得多的模型来猜，比如 LaMDA 家族 137B 的主模型，可以用 2B 或 8B 的小模型来猜。小模型的运算比大模型快得多，所以可以让它一口气往前走好几步，也就是算出接下来 m 个位置的 token 和对应的概率分布。
验证。由大模型对这 m 个 token 并行验算。这里不需要再一个一个计算：对于第 x+2 到第 x+m 个位置，大模型直接把小模型猜出的第 x+1 到第 x+m−1 个 token 当作上文，一次前向就同时得到每个位置的预测。大模型只读一遍权重，就完成了 m 个位置的计算，这正是图 1 里"多处理几个 token 几乎不花额外时间"的用武之地。

一次前向，同时验证所有位置

图 4点上方按钮或下方任一输出格，看它能"看到"哪些输入。Transformer 的因果注意力让每个位置只看自己和左边的 token，所以一次前向就能同时给出 m+1 个位置的预测，最后一格是全部猜对时白送的。

把猜测和验证连起来跑，效果是这样的：

小模型猜的大模型认可被拒绝大模型补上

普通解码一次一个

投机解码m = 4

两条赛道写的是同一句话，按同一个时钟走。

图 5时间条里每段蓝色是一次大模型前向，黄色是一次小模型前向（这里按大模型的 0.12 倍计）。

未完待续

下一部分会讲：大模型怎么判断小模型猜得对不对，猜错了怎么补，以及为什么这样做完全不会改变大模型的输出分布。想先看完整过程，可以读投机解码，可交互图解。

参考

Leviathan, Kalman, Matias. Fast Inference from Transformers via Speculative Decoding. ICML 2023.