什么是投机解码？（70分钟阅读）

# 投机解码：猜得好的艺术 Source: https://astralord.github.io/posts/speculative-decoding-the-art-of-guessing-well/ > 大语言模型的大部分解码时间并非花在计算上，而是花在等待上：每生成一个新 token，都要从内存中读取全部模型权重，却只是将它们与一个向量相乘。如果模型能以一次的代价验证多个 token，会怎样？ > > 在本文中，我们将从第一性原理推导投机解码，构建一个涵盖所有现代草稿器的通用框架，并逐一介绍已成为行业标准的架构——从 EAGLE-3 和多 token 预测，到 DFlash 和 DSpark 等块扩散式草稿器。最后我们将探讨如何训练草稿器，以及为什么最小化 KL 散度并不完全是我们想要的。 > > 本文是《Transformer 推理优化工具集》(https://astralord.github.io/posts/transformer-inference-optimization-toolset/) 的续篇。我们将沿用其符号约定、内存带宽算术和 KV 缓存计算，因此如果"解码是内存瓶颈"这句话听起来陌生，值得先浏览那里的前几节。 > > 部分投机解码技术将被省略。我们聚焦于**挂在目标模型上的草稿器**：在同一 GPU 上、读取目标模型隐状态并并行运行的小型网络。独立的小型 LLM 仅作为基线出现，通过字符串匹配（n-gram、提示词查找、后缀树）进行草稿的方法归入脚注1 (https://astralord.github.io/posts/speculative-deoding-the-art-of-guessing-well/#fn:JAC)，服务系统仅在影响草稿器构建和使用方式时才提及。 ## 为什么解码很慢 让我们回顾一下[文本生成](https://astralord.github.io/posts/transformer-inference-optimization-toolset/#kv-cache)分为两个阶段。在 prefill 阶段，模型一次性读入整个提示词，此时是计算瓶颈。在解码阶段，模型每次前向传播只生成一个 token，每次前向传播的[算术强度](https://astralord.github.io/posts/transformer-inference-optimization-toolset/#arithmetic-intensity-vs-opsbyte)远低于 GPU 的 `ops:byte` 比率。 在 batch size 为 1 时，每个线性层将其权重矩阵与单个 token 的隐状态（大小为 $d$ 的向量）相乘，因此每个权重从 HBM 读出后，仅用于一次乘法和一次加法，然后闲置到下一个步骤。 一个真实世界的例子：在 H100 GPU 上，以 bf16 格式运行一个 8B 模型。仅读取一次权重就需要： $$\underset{\text{weights, bf16}}{16 \text{ GB}} \ / \ \underset{\text{H100 HBM bandwidth}}{3.35 \text{ TB/s}} \approx 4.8 \text{ ms} \ \Rightarrow \ \le 208 \ \tfrac{\text{tokens}}{\text{s}} \text{ for an 8B model at batch size 1},$$ 无论 Tensor Core 多快都一样。与此同时，一个 token 的计算量是每个权重一次乘加，即 $2 \cdot 8 \cdot 10^9 = 16$ GFLOPs（忽略对 KV 缓存的注意力），H100 在 $989$ TFLOPs 的 dense bf16 算力下约需 $0.016$ ms 完成。 关键来了：如果我们给模型喂 8 个 token 而不是 1 个，计算量增加 8 倍，但付出的时间成本仍接近 4.8 ms，因为权重无论如何只读一次。该前向传播仍处于内存瓶颈状态，直到被评分的 token 数量达到 `ops:byte` 比率 $989 / 3.35 \approx 295$。 []()**8B 模型在 bf16 下于 H100 上的单次前向传播耗时与该次传播中被评分的 token 数 $n$ 的关系**：$\max \big( \frac{16 \text{ GB}}{3.35 \text{ TB/s}}, \frac{2 \cdot 8 \cdot 10^9 \cdot n}{989 \text{ TFLOPs}} \big)$，忽略 KV 缓存读取。直到 $n \approx 295$ 之前，前向传播受权重读取限制，因此额外的 token 几乎是免费的。* 批量处理请求是弥补这一差距的一种方式，服务引擎会积极地进行批处理。但批处理提高的是吞吐量，而非单个请求的延迟：每个用户仍然每 4.8 ms 才得到一个 token。要让单条序列加速，前向传播中额外的 token 必须来自同一条序列——而我们还不知道它们是什么。除非我们去猜测。 ### 草稿与验证 假设我们有一个廉价的**草稿器** $q$，它猜测下一个 token 的速度远快于我们实际想要采样的**目标**模型 $p$。让我们从贪心解码开始，即目标模型总是选择其最可能的 token。一轮**投机解码**如下所示： 1. 草稿器提议 $\gamma$ 个 token $x_{t+1}, \dots, x_{t+\gamma}$，延续已验证的上下文 $x_{\leq t}$。 2. 目标模型对扩展了全部 $\gamma$ 个猜测的上下文运行*一次*前向传播。借助因果注意力，这单次传播给出 $\gamma + 1$ 个位置上各自的下一个 token 预测。 3. 我们保留猜测中目标模型的 argmax 与草稿器一致的最长前缀。在第一个不一致处，我们取目标模型自身的 token；如果全部 $\gamma$ 个猜测都一致，则目标模型在最后一个猜测之后的预测便是免费获得的。 我们保留的每个 token 恰好是贪心解码本会产生的 token，因此输出不变，只是到达得更快。每一轮消耗一次目标模型前向传播，至少产出一个 token，最多产出 $\gamma + 1$ 个。已接受 token 的 KV 缓存条目在验证传播期间已写入，被拒绝的 token 的条目则直接丢弃。 []()**$\gamma = 4$ 的一轮贪心解码。** 草稿器逐个提议 token（浅蓝），然后目标模型在单次传播中对它们全部评分（橙色）。最长的一致前缀被接受（绿色），其余被拒绝（红色），并追加目标模型在第一个不一致处自身的 token：花一次目标模型前向传播的代价换来 3 个新 token。* ### 投机采样 当我们使用 temperature 进行采样时，仅靠贪心匹配是不够的：目标模型现在拥有一个分布 $p$ 而非唯一正确答案，我们希望输出精确遵循该分布。**投机采样**（Leviathan et al., 2023 (https://arxiv.org/pdf/2211.17192)；Chen et al., 2023 (https://arxiv.org/pdf/2302.01318)）将匹配规则推广为： $$\text{accept } x \sim q \ \text{ w.p. } \min\Big(1, \frac{p(x)}{q(x)}\Big), \qquad \text{otherwise sample from } \ \frac{(p - q)_+}{\sum_x (p(x) - q(x))_+},$$ 其中 $(\cdot)_+ = \max(\cdot, 0)$。直观地说，草稿器低估的 token（$q(x) \leq p(x)$）总是被保留，高估的 token 被稀释，残差分布恰好补足草稿器缺失的概率质量。 草稿 token 被接受的频率有多高？对草稿器自身的样本取接受概率的期望，得到**接受率**： $$\alpha = \mathbb{E}_{x \sim q}\Big[\min\Big(1, \frac{p(x)}{q(x)}\Big)\Big] = \sum_x \min\big(p(x), q(x)\big),$$ 即两个分布的重叠部分：两个模型放在同一 token 上的概率质量被接受，草稿器放在目标模型不放的地方的质量被浪费，而它遗漏的质量需要由残差来补充。缺失的部分有一个标准名称：即两个分布的**全变差距离**（total variation distance），它们能为任何 token 集合赋予的最大概率差： $$1 - \alpha = \operatorname{TV}(p, q) = \frac{1}{2}\sum_x \vert p(x) - q(x) \vert.$$ Leviathan et al. 将其记为 $D_{LK}$。草稿器在 TV 上越接近目标模型，其猜测存活的比例越高。 现在我们可以理解为什么输出仍然遵循 $p$。残差的归一化因子是 $\sum_x (p(x) - q(x))_+ = 1 - \alpha$（即被拒绝的概率），因此输出 token $x$ 的概率为： $$\begin{aligned}\Pr(\text{output} = x) &= q(x)\min\Big(1, \frac{p(x)}{q(x)}\Big) + (1 - \alpha)\,\frac{(p(x) - q(x))_+}{1 - \alpha} \\ &= \min\big(p(x), q(x)\big) + \big(p(x) - q(x)\big)_+ \\ &= p(x).\end{aligned}$$ 因此无论草稿器做什么，输出都精确遵循目标模型，而一个糟糕的草稿器只会让我们变慢。贪心匹配是 $p$ 和 $q$ 均为 one-hot 时的特例。对于 $\gamma$ 个 token 的草稿，该规则从左到右应用，在第一个被拒绝处停止。以下是用 JAX 实现的完整验证步骤。当全部 $\gamma$ 个 token 被接受时，下一个位置没有草稿器分布，将其视为零便可使残差变为目标模型自身的分布，用于 bonus token： ```python @jit def verify(key, p, q, draft): """p: [γ+1, V] target probs, q: [γ, V] drafter probs, draft: [γ] tokens.""" gamma = draft.shape[0] key_u, key_r = jax.random.split(key) p_x = p[jnp.arange(gamma), draft] q_x = q[jnp.arange(gamma), draft] accepted = jax.random.uniform(key_u, (gamma,)) < jnp.minimum(1.0, p_x / q_x) n = jnp.argmin(jnp.append(accepted, False)) # first rejection q_n = jnp.where(n < gamma, q[jnp.minimum(n, gamma - 1)], 0.0) # zero after γ residual = jnp.maximum(p[n] - q_n, 0.0) token = jax.random.categorical(key_r, jnp.log(residual)) return n, token # keep draft[:n], then append token ``` ### 共享噪声的采样 投机采样有一个不太方便的性质：其输出依赖于草稿器。在相同随机种子下，不同的草稿器会做出不同的接受/拒绝决策并采样不同的残差，因此文本会变化——尽管其分布不变。要看一个替代方案，让我们看看 LLM 的采样实际上是如何实现的。 对下一个 token 采样意味着从 $V$ 个 token 的词表上的 $\operatorname{softmax}(\ell)$ 中采样，其中 $\ell$ 是 logits。教科书方法对 logits 做归一化，计算概率的累积和，然后用一个均匀随机数进行逆变换。GPU 更偏好 **Gumbel-max 技巧**：对 logits 加上独立的 Gumbel 噪声，然后取 argmax： $$x = \arg\max_i \big(\ell_i + g_i\big), \qquad g_i = -\log(-\log u_i), \quad u_i \sim U(0, 1).$$ 为什么这能从 softmax 中采样？令 $\xi_i = e^{-g_i} = -\log u_i$，这是一个速率为 1 的指数随机变量。最大化 $\ell_i + g_i$ 等价于最小化 $\xi_i e^{-\ell_i}$，后者是以 $e^{\ell_i}$ 为速率的指数分布。想象独立的指数时钟：时钟 $i$ 最先响起的概率是其速率除以所有速率之和，因此： $$\Pr(x = i) = \Pr\Big(\xi_i e^{-\ell_i} < \min_{j \neq i} \xi_j e^{-\ell_j}\Big) = \frac{e^{\ell_i}}{\sum_j e^{\ell_j}} = \operatorname{softmax}(\ell)_i.$$$ 一次逐元素操作加一个 argmax 就替代了词表上的归一化和前缀和。例如，vLLM 的采样器运行的正是这场竞赛，对 $p_i / \xi_i$ 取 argmax，因为 `torch.multinomial` 会迫使 GPU 与 CPU 同步。 对我们而言，该技巧的重要性在于另一个原因：它将所有随机性与模型分离。一旦噪声 $g$ 固定，采样就变成了确定性的 argmax，就像贪心解码一样。所以让草稿器和目标模型共享噪声。在草稿位置 $k$，两者使用相同的向量 $g_k$：草稿器提议 $x_{t+k} = \arg\max_i \big(\log q_k(i) + g_{k,i}\big)$，目标模型在验证传播中计算自身的 $\arg\max_i \big(\log p_k(i) + g_{k,i}\big)$，我们保留两者一致的最长前缀，加上目标模型在第一个不一致处的 token。验证变成了精确匹配，如同贪心情形。 这就是 **Gumbel 耦合**（Daliri et al., 2024 (https://arxiv.org/pdf/2408.07978)），其输出正是目标模型用该噪声自行采样会得到的结果。它是**草稿器无关的**：固定种子时，每个草稿器产生相同的文本，只是快慢不同——这使投机解码可复现，且更容易测试和调试。 []()**在 "The cat sat on the" 之后 token 的 Gumbel-max 采样。** 每行是一个候选 token：橙色点是目标模型的分数 $\log p + g$，蓝色点是草稿器的分数。最右侧的橙色点是目标模型的采样，最右侧的蓝色点是草稿，当两者重合时草稿被接受。共享噪声时，同一行的两个点一起移动，两个 argmax 大多数时候一致；使用独立噪声时，两者很少一致。点击标题切换，并生成新噪声以观察匹配率收敛。* 代价是更低的接受率。投机采样能达到最佳的可能一致性 $\alpha = 1 - \operatorname{TV}(p, q)$，因为验证器能看到两个分布。仅共享噪声的两个模型的一致概率为： $$\Pr\big[\text{match}\big] = \sum_j \frac{1}{\sum_i \max\big(\frac{p(i)}{p(j)}, \frac{q(i)}{q(j)}\big)} \ \geq \ \frac{1 - \operatorname{TV}(p, q)}{1 + \operatorname{TV}(p, q)},$$ 在最坏情况下，任何不进行通信的方案都无法保证超过右侧的下界。共享噪声正是使之生效的原因：使用独立噪声时，两个采样一致的概率仅为 $\sum_i p(i)\, q(i)$。在上述示例中，共享噪声时为 0.679，独立噪声时为 0.181，而投机采样为 $\alpha = 0.70$；在真实 LLM 分布上，Gumbel 耦合与 $1 - \operatorname{TV}$ 之间的差距通常同样很小。 在 JAX 中，对于给定草稿路径上的分布，整个草稿-验证步骤只需在同一噪声上进行几次 argmax： ```python @jit def gumbel_verify(key, log_p, log_q): """log_p: [γ+1, V] target log-probs, log_q: [γ, V] drafter log-probs along the draft. The same key, i.e. the same noise, is used by the drafter and the target.""" g = jax.random.gumbel(key, log_p.shape) draft = jnp.argmax(log_q + g[:-1], axis=-1) # what the drafter proposes target = jnp.argmax(log_p + g, axis=-1) # what the target samples n = jnp.argmin(jnp.append(draft == target[:-1], False)) # first mismatch return draft, n, target[n] # keep draft[:n], then append target[n] ``` ### 我们能拿到多少 token？ 首先假设每个草稿 token 以相同的概率 $\alpha$ 独立地被接受。一轮在第一个被拒绝处停止，因此被接受的草稿 token 数服从截断在 $\gamma$ 处的几何分布，加上 bonus token，每轮产出的期望 token 数为： $$\mathbb{E}[\#\text{tokens}] = \frac{1 - \alpha^{\gamma + 1}}{1 - \alpha}.$$ 实际上，草稿器离已验证上下文越远猜测得越差，因此接受率取决于草稿中的位置。如果 $\alpha_k$ 是在前 $k-1$ 个 token 全部被接受的条件下，第 $k$ 个 token 被接受的平均概率，那么**接受长度** $\tau$（即每轮产出的期望 token 数）为： $$\tau \approx 1 + \sum_{k=1}^{\gamma} \prod_{i=1}^{k} \alpha_i.$$ 这是论文中最常报告的数字，草稿器之间就在此指标上竞争。 现在来算实际耗时。