What I Gave Up to Make Speculative Decoding Faster
本文为原文前 6,000 字符的节选翻译，完整内容请查看原文。
What I Gave Up to Make Speculative Decoding Faster
在创意写作提示词上，我的目标大语言模型（LLM）拒绝了较小草稿模型提出的大多数标记。推测解码的速度比单独运行目标模型还要慢。我希望能挽回这些浪费的工作。如果我让草稿模型对输出产生更大的影响，生成速度能提高多少？无损推测采样保留了目标模型的输出分布。放宽这一要求可能会提高接受率，但由此产生的文本需要更仔细的审视。我在我的 macOS Rust 推理引擎 mini-vllm-rs 中通过“指导解码”（mentored decoding）探索了这种权衡。在 Apple M2 上，使用量化的 Qwen2.5 7B 作为目标模型，0.5B 作为草稿模型，我最终找到了一种更快的配置。至于何时使用它，则成了另一个问题。
A knob for accepting more draft tokens 在早期的基准测试中，代码重构提示词在贪婪解码下表现良好。一种自适应调整草稿标记数量的策略通常会选择 8 个标记的建议，实现 98.5% 的接受率，并将解码吞吐量比仅使用目标模型提高了 38.3%。然而，在创意写作提示词上，同样的策略在贪婪解码下产生的解码吞吐量比仅使用目标模型低 36.8%。随后，我使用温度为 0.7 的多项式采样（采样基准）探索了创意写作提示词。固定草稿标记数量为 4 的无损推测采样仅实现了 28.6% 的接受率。
指导解码允许从目标分布向草稿分布进行有界偏移。 我实现了它的线性插值构造。设 $p$ 为目标分布，$q$ 为草稿分布，两者均在温度和重复惩罚之后。它们的总变差距离为： $d=\operatorname{TV}(p,q)=\frac12\sum_x |p(x)-q(x)|$ 给定一个 TV 预算 $D$，插值权重为： $\theta=\min(1,D/d)$ 如果 $d=0$，则分布完全匹配，我们取 $\theta=1$。由此产生的输出分布为： $\pi=(1-\theta)p+\theta q$ 从 $q$ 中采样的草稿标记 $x$ 以以下概率被接受： $a(x)=\min\left(1,\frac{\pi(x)}{q(x)}\right)$ 拒绝后，采样器从归一化的正残差 $\max(p-q,0)$ 中抽取替换项。总之，接受和修正共同产生了分布 $\pi$。在 TV 预算为 0 时，这保留了目标分布。在 TV 预算为 1 时，预算涵盖了所有可能的 TV 距离，因此 $\theta=1$ 且 $\pi=q$。此时，每个草稿标记都会被接受。中间预算限制了每一步相对于目标的分布偏差。它们并不衡量输出损失了多少语义质量。
Even full acceptance did not buy a speedup 我首先将 TV 预算设为 1，因此每个草稿标记都被接受。即便如此，固定 4 个标记的生成耗时 24.425 秒，而仅使用目标模型耗时 20.971 秒。拒绝虽然消失了，但执行成本依然存在。每一轮仍然自回归地生成 4 个草稿标记，用目标模型验证它们的位置，进行采样，并更新两个模型的缓存状态。我的实现不会在完全接受的建议后发出额外的目标奖励标记，因此一轮成功的生成仅产生 4 个输出标记。仅使用草稿模型的对照组耗时 3.372 秒。在 TV=1 时，目标验证是遵循草稿分布的一种昂贵方式。该端点有助于诊断剩余成本，但我希望在保留更多目标分布的同时实现加速。
Finding the break-even point 为了找到推测何时能胜过仅使用目标模型的端到端延迟，我在完全接受的情况下测试了 6 和 8 个固定草稿标记数量。然后，我降低了 TV 预算，以找到仍能产生明显加速的最小测试值（基准报告）。提示词和温度保持不变。每个请求生成 256 个标记，忽略 EOS。预热后，每个设置进行了 3 次独立的未追踪计时调用，每次覆盖 3 个随机种子。
| 配置 | 接受率 | 平均每请求轮数 | 中位数端到端延迟 |
|---|---|---|---|
| 仅目标模型 | — | — | 20.971 s |
| 固定 4, TV=1 | 100% | 64 | 24.425 s |
| 固定 6, TV=1 | 100% | 43 | 20.587 s |
| 固定 8, TV=1 | 100% | 32 | 18.508 s |
| 固定 8, TV=0.75 | 93.4% | 34 | 18.820 s |
| 固定 8, TV=0.5 | 63.5% | 46 | 22.900 s |
在观察到的变异范围内，固定 6 与仅使用目标模型大致相等。固定 8 明显跨越了基准线。在完全接受的情况下，它每个请求需要 32 轮验证，而不是固定 4 的 64 轮，从而允许更多的输出共享每次验证调用。验证也跨越了内核调度边界，正如我在连续批处理实验中所做的那样。在这里，4 行验证使用了行式 GEMV，而完整的 6 行和 8 行建议使用了常规的量化矩阵乘法。建议长度和内核调度同时改变，因此这些测量结果并未隔离它们各自的贡献。
What did I get in exchange? 对于固定 8，将预算从 0.5 增加到 0.75，平均验证轮数从 46 轮减少到 34 轮。生成速度从比仅使用目标模型慢 9.2% 变为快 10.3%。在 TV=0.75 时，每个种子的中位数提高了约 8–10%。为了判断这种加速换来了什么，我检查了创意写作的输出。提示词将机械师玛拉（Mara）和档案管理员伊利安（Ilyan）置于一个废弃 20 年的月球温室中调查求救信号。它指定了尽管缺乏电力仍有新鲜冷凝水，以及从地球反射的蓝色光斑。在一个仅使用目标模型的样本中，角色将冷凝水视为有人可能来过或某些事物正在改变的证据。在使用相同种子的 TV=0.75 输出中，屋顶显示：“一缕细微的蓝光射入——这是地球的红外线，表明这并非该系统的终结，而只是其正常运行的中断。”红外线并非可见的蓝光。这段文字还从光线直接跳跃到对系统的诊断，而没有解释其中的联系。在 TV=1 时，同一个种子产生了身份错误。Immed