Speculative Decoding：用小模型帮大模型提前草拟答案

前面几篇文章里，我们优化 LLM 推理的方向主要集中在系统层。

KV cache 避免重复计算历史 token。

Continuous batching 提高多请求 decode 的 GPU 利用率。

PagedAttention 改善 KV cache 的显存管理。

Prefix caching 减少重复 prefill。

Chunked prefill 避免长 prompt 阻塞 decode。

这些优化都很重要，但它们大多没有改变一个基本事实：大模型在 decode 阶段通常还是一轮生成一个 token。

也就是说，如果一个请求要生成 512 个 token，大模型通常需要执行 512 次 decode forward。

Speculative decoding 想解决的是另一个问题：

能不能让大模型一次确认多个 token，而不是每次只生成一个？

它的核心思路是，引入一个更小、更快的 draft model。小模型先根据当前上下文草拟多个 token，然后大模型一次性验证这些 token 是否可以接受。如果草拟 token 被接受，大模型就相当于一次推进了多个 token，从而减少大模型 decode 的步数。

这个方法很有意思，因为它不是简单地换一个更快的 kernel，也不是优化 KV cache 存储，而是改变 decode 的执行模式。

不过，它也不是免费的。

它会引入小模型的计算成本，会增加验证逻辑，会影响调度和 KV cache 管理，还会让请求状态更复杂。

这一篇我们先从直觉讲起，再讨论 speculative decoding 如何接入 mini vLLM。

为什么 decode 慢

在自回归生成里，每个 token 都依赖前面的 token。

模型生成 token t 之后，才能把它接回上下文，再生成 token t + 1。

所以最朴素的 decode 是串行的。

输入上下文
大模型生成 token_1

输入上下文 + token_1
大模型生成 token_2

输入上下文 + token_1 + token_2
大模型生成 token_3

即使有 KV cache，大模型每生成一个 token 仍然需要执行一次 forward。

KV cache 只是避免重新计算历史 token，但不能打破自回归依赖。

这就是 decode 的根本限制。

如果输出很长，大模型 decode 的轮数就很多。

假设大模型每次 decode 需要 20ms，生成 500 个 token 就需要大约 10 秒。实际系统中还会有 batch、调度、采样、网络输出等因素，但 decode 步数仍然是核心成本之一。

Speculative decoding 的出发点是：虽然大模型必须保证最终输出质量，但也许不需要让大模型亲自一步一步生成所有 token。

小模型可以先猜。

大模型负责检查。

如果小模型猜得对，大模型就能少跑几轮。

草拟和验证

Speculative decoding 可以分成两个阶段。

第一阶段是 draft。

小模型基于当前上下文，连续生成多个候选 token。

例如一次草拟 4 个 token：

draft model 生成:
d1, d2, d3, d4

第二阶段是 verify。

大模型不是只验证 d1，而是一次性处理这 4 个 token，并计算每个位置的大模型分布。

也就是说，大模型会看：

上下文 + d1 + d2 + d3 + d4

然后判断这些 draft tokens 中有多少可以接受。

如果 d1、d2、d3 都被接受，d4 被拒绝，那么最终输出可以一次追加 d1、d2、d3，然后由大模型在拒绝位置重新采样一个 token。

这样，大模型一次 forward 推进了多个 token。

如果小模型草拟质量很高，接受率高，那么大模型 decode 次数会明显减少。

如果小模型草拟质量很差，大部分 token 都被拒绝，那么 speculative decoding 可能反而变慢，因为额外跑了小模型。

所以 speculative decoding 的收益取决于两个关键因素：

draft model 是否足够快
draft tokens 是否有较高接受率

小模型太慢，不划算。

小模型太差，也不划算。

理想的小模型应该足够快，同时生成分布和大模型比较接近。

一个直观例子

假设当前上下文是：

The capital of France is

小模型草拟了 4 个 token：

Paris, and it

大模型验证后认为这些 token 都合理，于是全部接受。

那么大模型一次 verify，就让输出前进了 4 个 token。

如果没有 speculative decoding，大模型需要 4 次 decode forward。

生成 Paris
生成 ,
生成 and
生成 it

现在变成一次 verify。

当然，verify 的输入长度是 4 个 token，不是 1 个 token，所以这一次 forward 比普通单 token decode 更重。

但它通常仍然比大模型连续跑 4 次 decode 更划算。

因为大模型可以在一次 forward 中并行处理 draft tokens。

这就是 speculative decoding 的性能来源。

它用一次较大的验证 forward，替代多次小的 decode forward。

接受和拒绝不是简单 argmax 对比

很多人第一次理解 speculative decoding 时，会以为验证就是比较小模型和大模型的 argmax 是否一致。

如果小模型猜的 token 正好也是大模型最想选的 token，就接受；否则拒绝。

这种理解很直观，但不够完整。

在采样场景下，目标不是让输出等于大模型 greedy 的结果，而是让最终采样分布和直接从大模型采样一致。

经典 speculative decoding 会基于 draft model 分布和 target model 分布做接受拒绝采样。

简单说，小模型提出 token，大模型给出该 token 在自己分布下的概率，然后根据一定规则决定是否接受。

如果拒绝，则需要从修正后的分布中重新采样，保证最终结果仍然符合大模型分布。

这里的数学细节会比 greedy 复杂。

为了 mini vLLM 的教学版本，我们可以先从 greedy speculative decoding 开始。

也就是：

draft model 草拟若干 token
target model 验证这些 token
如果 draft token 等于 target greedy token，则接受
遇到第一个不一致就停止

这个版本不覆盖完整采样语义，但非常适合理解系统流程。

等系统结构跑通后，再扩展到概率接受版本。

这符合我们系列一贯的路线：先正确理解主流程，再逐步处理完整复杂度。

Greedy 版本的 speculative decoding

先看 greedy 版本。

假设 draft model 一次生成 4 个 token：

draft tokens:
d1, d2, d3, d4

target model 读取当前上下文加这些 draft tokens。

它会输出每个位置的 logits。

我们取 target model 在对应位置的 greedy token。

target tokens:
t1, t2, t3, t4

然后从左到右比较。

如果：

d1 == t1
d2 == t2
d3 != t3

那么接受 d1 和 d2。

在 d3 的位置，使用 target model 的 t3 作为正式输出。

所以这一轮最终追加：

d1, d2, t3

后面的 d4 丢弃。

为什么要在拒绝位置追加 target token？

因为这一轮 target model 已经计算出了这个位置的结果。如果不利用它，就浪费了。

这样，一次 speculative step 至少能追加一个 token。最好的情况下追加 draft_len 个 token，甚至在某些实现中可以追加 draft_len + 1 个 token。

教学版可以先采用简单语义：接受连续匹配的 draft tokens，遇到第一个不匹配时追加 target token，然后结束本轮。

这个版本很好理解，也方便测试。

Speculative decoding 和 KV cache

Speculative decoding 接入推理引擎后，最麻烦的部分之一是 KV cache。

因为现在一次 step 可能不再只生成一个 token，而是生成多个 token。

对于 draft model，它需要自己的 KV cache。

对于 target model，它也需要自己的 KV cache。

这两份 cache 不能混用，因为两个模型参数不同，hidden states 不同，K/V 当然也不同。

所以如果使用一个小模型和一个大模型，就需要两套 KV cache 管理。

draft KV cache
target KV cache

这会增加显存占用。

如果 draft model 很小，它的 KV cache 也会小一些，但仍然不是零。

在教学版里，可以先把 draft model 的 cache 管理做得简单一点，甚至先不接入 paged KV cache，只用框架自己的 past_key_values。

但如果要做更完整的 speculative serving，就要让 draft model 也进入类似的 cache 管理体系。

更复杂的是验证阶段。

target model 一次验证多个 draft tokens。这意味着它要把这些 tokens 的 K/V 写入 target KV cache。

但如果后面某些 draft tokens 被拒绝，它们对应的 target KV cache 不能保留。

比如 draft_len 是 4，target model 计算了 d1、d2、d3、d4 的 K/V。

最后只接受 d1、d2，并在 d3 位置改用 target token t3。

那么 d3、d4 的 KV cache 里原本写入的是 draft tokens 的 K/V，不一定对应最终序列。

这些 cache 必须回滚或重写。

这就是 speculative decoding 和 KV cache 管理的复杂点。

一个简单教学实现可以先避免复杂回滚：验证阶段先不要直接提交写入正式 KV cache，而是使用临时 cache。确认接受哪些 tokens 后，再把正式接受的 tokens 写入 target cache。

这种做法性能不高，但语义清楚。

更高性能的实现会设计更精细的 cache commit 和 rollback 机制。

为什么需要 commit / rollback 语义

在普通 decode 中，每轮只生成一个 token。

模型输入 last token，写入这个 token 的 KV，采样 next token。

这条路径很简单，写入基本可以立即提交。

Speculative decoding 不一样。

target model 会先基于 draft tokens 做验证，但这些 draft tokens 可能不会全部进入最终输出。

所以它的写入应该分成两个阶段：

prepare:
  为 draft tokens 计算可能的 KV

commit:
  只提交最终接受的 tokens

rollback:
  丢弃未接受 tokens 对应的临时 KV

如果没有这种语义，Block Manager 会很难保证 block table 正确。

因为 Sequence 的最终上下文只能包含被接受的 tokens。

未接受 token 的 KV 不能出现在 block table 中。

否则后续 PagedAttention 会读到错误历史。

所以 speculative decoding 会迫使我们把 KV cache 生命周期设计得更精细。

对于 mini vLLM 的第一版，可以把 speculative decoding 做成一个独立路径，先用简单 cache 实现，重点理解算法流程。

等后面系统成熟，再把它和 Block Manager 深度融合。

Scheduler 视角：一次 step 不再等于一个 token

引入 speculative decoding 后，Scheduler 的成本模型也会变化。

以前 running Sequence 每轮 decode 大致消耗一个 token 的计算预算，并产生一个 token。

现在一个 running Sequence 可能会先运行 draft model 生成多个 token，再运行 target model 验证多个 token，最后追加多个 token。

这意味着：

一轮调度的计算成本增加
一轮调度的输出 token 数也可能增加

Scheduler 不能再简单地认为 decode sequence 的 token 成本等于 1。

它要知道 speculative decoding 的参数，例如 draft_len。

如果 draft_len 是 4，那么一次 speculative step 可能包含：

draft model 4 次小 decode
target model 1 次多 token verify

或者 draft model 自身也可以用某种批量方式生成 4 个 tokens。

这会影响调度预算。

尤其是在 batch 中，不同请求可能有的启用 speculative decoding，有的没有启用。

这让调度更加复杂。

教学版 mini vLLM 可以先做一个简单限制：要么整个 engine 启用 speculative decoding，要么全部不用。这样调度器不用在同一 batch 里混合两种 decode 模式。

后续再支持 per request 配置。

这种限制是合理的。因为我们现在的目标是讲清楚机制，而不是一开始就覆盖所有生产细节。

Speculative decoding 的第一版接口

从工程结构上看，可以把 speculative decoding 设计成一个新的 decoding backend。

普通 decode 是：

输入 last token
target model forward
sampler 采样 next token
append 一个 token

Speculative decode 是：

draft model 基于当前上下文生成 draft tokens
target model 验证 draft tokens
accept 若干 tokens
append accepted tokens

所以可以抽象一个 Decoder 接口。

class Decoder:
    def decode(sequence) -> DecodeOutput:
        ...

普通 decoder 返回一个 token。

Speculative decoder 返回一个或多个 token。

class DecodeOutput:
    sequence_id: str
    token_ids: list[int]
    accepted_count: int

这样 Engine 不需要关心底层是普通 decode 还是 speculative decode。

它只需要把返回的 token_ids 追加到 Sequence。

不过，追加多个 token 时，Sequence 的状态更新也要支持批量。

sequence.append_tokens(token_ids)

同时要注意 stop condition。

如果一次接受了多个 token，其中间某个 token 命中了 EOS 或 stop words，那么后面的 tokens 不应该继续输出。

所以 append_tokens 不能只是简单 extend list。

它需要逐个检查停止条件。

这也是 speculative decoding 带来的一个细节：一次 step 产生多个 token，但请求可能在中间某个 token 就结束。

和 Sampling 的关系

Speculative decoding 和 sampling 系统关系非常紧密。

如果我们只做 greedy speculative decoding，事情比较简单。

draft model 用 greedy 生成 draft tokens。

target model 用 greedy 验证。

如果一致就接受，不一致就用 target greedy token。

但如果支持 temperature、top p、top k，就复杂很多。

因为最终输出必须保持和 target model 的采样分布一致。

这时不能简单比较 argmax。

需要使用接受拒绝采样。

概念上，小模型提出 token x。

draft model 给出概率 q(x)。

target model 给出概率 p(x)。

接受概率和 p(x)、q(x) 有关。

如果拒绝，需要从修正后的 target 分布中重新采样。

这里的数学细节比较多，而且实现上要拿到 draft 和 target 两边的 token probabilities。

所以第一版建议只支持 greedy speculative decoding。

或者明确说明：当前 speculative decoding 只用于 temperature = 0 的场景。

这样既能讲清系统流程，也能避免采样语义不严谨。

后续如果要写扩展篇，可以单独讲完整 speculative sampling。

这比在第一版里混在一起更清晰。

接受率决定收益

Speculative decoding 是否有收益，关键看接受率。

如果 draft model 提出的 tokens 大部分都被 target model 接受，那么 target model 每次 verify 可以推进多个 token，收益明显。

如果接受率很低，每次只接受 0 或 1 个 token，那就不划算。

因为系统额外运行了 draft model，还增加了验证和状态管理开销。

可以定义一个指标：

acceptance_rate = accepted_draft_tokens / total_draft_tokens

还可以记录每轮平均接受 token 数。

average accepted tokens per speculative step

如果 draft_len 是 4，但平均只接受 1.1 个 token，收益可能很有限。

如果平均接受 3 个 token，就比较有价值。

接受率受很多因素影响。

draft model 和 target model 越接近，接受率越高。

任务越确定，接受率越高。

temperature 越高，采样越随机，接受率通常越难稳定。

代码补全、格式化输出、模板化回答，可能比较适合 speculative decoding。

高度开放的创意写作，draft model 可能更容易和 target model 分歧。

所以 speculative decoding 的 benchmark 必须记录接受率。只看 tokens/s 不够。

性能收益不是必然的

Speculative decoding 很吸引人，但它不是免费加速。

它可能变快，也可能变慢。

收益来自减少 target model decode 次数。

成本来自新增 draft model 计算、target model 多 token verify、KV cache 管理、接受拒绝逻辑和调度复杂度。

可以粗略理解为：

如果小模型草拟成本 + 大模型验证成本
小于
大模型逐 token decode 成本
那么有收益

但这个比较不是固定的。

batch size、draft_len、接受率、模型大小、硬件利用率都会影响结果。

如果 target model 原本 decode batch 已经很大，GPU 利用率很高，那么 speculative decoding 的相对收益可能降低。

如果 target model decode batch 小，单 token decode 很低效，那么 speculative decoding 可能更有帮助。

如果 draft model 运行在同一张 GPU 上，它会和 target model 抢资源。

如果 draft model 可以运行在另一张 GPU 或更便宜的硬件上，收益模型又不同。

所以 benchmark speculative decoding 时，不能只测一个场景。

要分别看：

不同 draft_len
不同接受率
不同 batch size
短输出和长输出
greedy 和 sampling
draft model 与 target model 的大小比例

这也是为什么它适合作为一个高级篇，而不是早期核心篇。

它很有价值，但系统复杂度也高。

mini vLLM 的第一版 speculative decoding

对于 mini vLLM，我建议第一版这样做。

只支持 greedy。

只支持 running decode 阶段。

不处理 prefill speculative。

draft model 使用一个小模型，可以先用框架自己的 generate 或手写 greedy decode 得到 draft tokens。

target model 用当前 mini vLLM 的路径验证 draft tokens。

KV cache 先采用简单安全实现，避免复杂 rollback。

可以先不把 speculative decoding 和 PagedAttention 深度融合，而是在正确理解流程后，再考虑优化。

第一版关注这些指标：

draft_len
accepted_tokens
acceptance_rate
target_model_forward_count
output_tokens_per_second

和普通 greedy decode 对比时，重点看 target model forward 次数是否下降。

如果生成 100 个 tokens，普通 decode 需要大约 100 次 target forward。

Speculative decoding 如果平均每轮接受 3 个 token，则 target forward 次数可能降到 30 多次。

当然，总耗时还要加上 draft model 的成本。

这个实验能很好地展示 speculative decoding 的基本价值。

小结

这一篇我们讨论了 speculative decoding。

它和前面的优化方向不同。KV cache、PagedAttention、continuous batching 主要是在优化大模型每一步 decode 的执行效率和资源管理。Speculative decoding 则试图减少大模型 decode 的步数。

它的核心思想是：

draft model 先草拟多个 tokens
target model 一次性验证这些 tokens
被接受的 tokens 直接追加到输出
被拒绝的位置由 target model 接管

如果 draft model 足够快，且草拟 token 的接受率足够高，那么 target model 可以用更少的 forward 生成同样数量的 tokens，从而提升吞吐或降低延迟。

但 speculative decoding 也会引入明显复杂度。

它需要 draft model，需要验证逻辑，需要处理接受拒绝，需要维护 target 和 draft 的 KV cache，还可能需要 commit / rollback 语义。对于 sampling 场景，还要保证最终分布和 target model 一致，不能简单比较 argmax。

因此，mini vLLM 的第一版可以先实现 greedy speculative decoding，把重点放在系统流程上。等主链路正确后，再扩展到完整的概率接受版本。

下一篇文章，我们会进入多 GPU 推理，讨论 tensor parallel。

到目前为止，我们默认模型可以放进一张 GPU。但更大的模型往往无法单卡容纳。Tensor parallel 的目标是把模型权重切到多张 GPU 上，让多个设备共同完成一次 forward。