vLLM 源码完全拆解（十）：投机解码——草稿模型、批量验证与 token 接受流水线

从「每一步一个 token」到「一个 batch 的草稿」#

标准自回归解码是这样的：给定已生成的序列，模型前向一次，输出下一个 token 的概率分布，采样出 $y_t$，把它接到序列末尾，再前向一次…如此串行往复。每步之间严格依赖，无法并行。对 Transformer 而言，生成第 $t$ 个 token 要重新算一遍整条序列的自注意力，计算量主要花在已经算过的 token 上——这就是解码阶段的算术强度（每 token 浮点 / 每 token 访存）很低、高度受限于 KV Cache 内存带宽的原因。

投机解码的思路（DeepMind 2022 那篇把它形式化）很反直觉：先让一个小得多的草稿模型（draft model）快速自回归地猜 $K$ 个 token，组成一个长度为 $K$ 的草稿序列；然后把这个草稿塞给目标模型，让目标模型一次前向，给出这 $K$ 个位置各自的预测分布；最后从后往前逐 token 检查，能接受的保留、不能接受的从那个位置重新采样。 如图：

图里左上是一个小目标模型循环生成草稿 token，右侧是大目标模型对整条草稿做一次并行前向，下面是从后往前做「接受 / 拒绝」的验证。这张图来自 DeepMind 的那篇原始论文 Fast Inference from Transformers via Speculative Decoding。

为什么它能「无损」加速#

关键在于接受准则。设草稿模型在某个位置给 token $x$ 的概率是 $q(x)$，目标模型给的概率是 $p(x)$。我们抛一个均匀随机数 $u \in [0,1)$，满足

就接受这个草稿 token；否则拒绝，并从调整后的分布

中重新采样一个 token 作为真正的输出。这个 $\tilde p$ 就是「目标分布减去草稿分布在 $x$ 处的那部分质量，再重新归一化」——直观上，如果草稿模型对某个 token 太有信心（$q$ 大），但目标模型并没那么认可（$p$ 小），那么这个 token 被接受的概率 $\frac{p}{q}$ 就低；被拒绝后，我们把它「按草稿的过分自信量」剔除，再从剩余质量里采样，从而保证重采样后的边际分布恰恰等于目标分布 $p$。这就是「无偏」（unbiased）的来源：只要草稿分布 $q$ 与目标分布 $p$ 一致或更激进，接受后得到的 token 分布与直接采样目标模型完全等价，代价只是多算了草稿的那几次便宜前向。

vLLM 在文档里给过一个形象的比喻：目标模型像一位严谨的教授，草稿模型像一个抢答的学生。学生先小声报一串答案，教授听一遍，认可的就点接受、不认可的当场指出并给出正确答案。因为大多数 token 其实很容易猜（“the”、“of” 这类高频词），教授一次前向就能「盖章」通过好几个 token——省下了原本要串行做的几次昂贵前向。

算法一页纸，工程十层楼#

算法本身只有一张图（上面那张），但它在真实推理框架里落地时，几乎每一层都会遇到新问题：

1. 草稿模型要能随时接入：vLLM 得支持 EAGLE（依赖目标模型的特征/hidden states）、独立的 draft LLM（有自己独立的 tokenizer 甚至词表）、Medusa 这种挂在目标模型边上的 MLP 头、以及 ngram 这种根本不加载模型的「伪草稿」。四种方式输入输出都不一样。
2. 草稿 token 怎么进目标模型：草稿序列的 token 不能和普通请求一样处理，它要作为「待验证的候选」塞进目标模型的前向，得为它构造出并行验证所需的批展开（batch expansion）和注意力元数据。
3. 验证怎么保证无偏：拒绝采样需要 draft_probs 和 target_logits 两套分布，还要处理 greedy/随机采样、temperature/top-k/top-p 这些采样参数，外加一个「bonus token」的收尾处理。
4. 多步草稿怎么和 CUDA Graph 合作：草稿模型通常要一次生成 $K>1$ 个 token，这意味着它自己要多步自回归；而 CUDA Graph 要求固定的 batch shape，两者必须协同。
5. KV Cache 和调度器怎么办：草稿 token 临时占用 KV 块，接受/拒绝后要回收；调度器要能感知「这批请求处于投机解码模式」。
6. 限制：不是所有模型/并行配置都能投机解码，比如张量并行度、流水线并行、词表异构等，各有约束。

这就是本篇正文的主体。下面我们从「谁在调谁」的总架构开始，层层往下拆。

主循环里插进来的一段#

GPUModelRunner._execute_model_step() 的伪代码骨架大致是：

1
1. 目标模型 prefill / decode 前向一次  →  target_logits
2
2. 目标模型先做一次"第 0 个 token"的采样   （本来的自回归采样）
3
3. if spec_config is set:
4
     a. 提取 target_hidden_states（EAGLE 要用的特征）
5
     b. proposer.prepare_next_token_ids_padded(...)   # 准备草稿输入
6
     c. proposer.propose(...)                          # 草稿模型生成 K 个候选 token
7
     d. proposer.prepare_next_token_ids_cpu(...)
8
     e. rejection_sampler(...)                          # 目标 logits 验证草稿
9
4. 处理输出，更新请求状态

注意第 2 步：目标模型本身仍然先生成一个 token。这个 token 是这条序列真正的、由 target logits 决定的输出（也是投机解码里「第 0 个接受 token」的基础），它同时作为草稿模型生成后续 $K$ 个 token 的「上下文起点」（EAGLE 里更是直接把它的 hidden state 喂给草稿层）。可以认为：投机解码不是「跳过目标模型」，而是「目标模型每跑一次，顺带验证 $K$ 个便宜草稿」。

Proposer：统一所有草稿方法的基类#

V1 用 SpecDecodeBaseProposer（vllm/v1/spec_decode/llm_base_proposer.py，约 1900 行）这一个类统一了口径。它的 method 字段决定了具体走哪条草稿路，代码里通过 isinstance(model, (...)) 或 model_type 分派。它的职责很重：持有一个草稿模型（或者根本不持有模型，如 ngram）、负责多步自回归生成草稿、构建草稿的注意力元数据、对接 CUDA Graph。

几个关键字段（都来自 SpecDecodeBaseProposer.__init__）：

不同草稿方法在 vLLM V1 里对应这些 Proposer（都在 vllm/v1/spec_decode/ 下）：

调度器与 KV：投机解码怎么融入既有框架#

前面第三、四篇讲过的 Scheduler 和 KVCacheManager 在这里并没有被推翻，只是多了一层标记。InputBatch / SamplingMetadata 里会带上投机相关的字段（如 num_speculative_tokens、spec_decode_meta），调度器给投机请求正常分配 KV 块，只是 draft token 会被作为「临时占用」对待：如果这轮草稿 token 被拒绝，它的 KV 位置（spec_decode_meta 里的 slot mapping）会被标记成 PADDING_SLOT_ID（值 -1），下一轮就不会再读它。这正好呼应我们对 KV Cache「有借有还」的理解——草稿的 KV 是可以作废的临时空间。

我们把「谁调谁」串成一张总图（基于前面九篇的框架）：

1
Scheduler.schedule ──> 产出本轮要执行的请求 + 分配 KV block
2
        │
3
        ▼
4
GPUModelRunner._execute_model_step
5
        │
6
        ├─ (1) 目标模型前向 + 第 0 个 token 采样
7
        │
8
        ├─ 若启用了投机　◄──── spec_config 存在
9
        │     ├─ 提取 target hidden states (EAGLE/MTP)
10
        │     ├─ proposer.prepare_next_token_ids_padded(...)
11
        │     ├─ proposer.propose(...)      ──> 草稿模型多步生成 K 个 token
12
        │     │        │
13
        │     │        └─ 内部: CUDA Graph 分派 + set_forward_context + 多步循环
14
        │     ├─ proposer.prepare_next_token_ids_cpu(...)
15
        │     └─ rejection_sampler(...)     ──> 用 target logits 验证草稿
16
        │
17
        └─ 输出 accepted tokens，更新请求与 KV

下面我们按这三段展开：草稿模型体系（propose 里到底谁在生成草稿）、验证与接受（rejection sampler 怎么做）、多步执行与 CUDA Graph（proposer 怎么把 K>1 跑起来）。

NgramProposer：不加载任何模型#

NgramProposer（vllm/v1/spec_decode/ngram_proposer.py）是投机解码里最「轻」的实现——它根本不加载草稿模型，而是在已生成的 token 串里找「历史中出现过的相同后缀」，把后缀后面紧跟的那段 token 当作草稿。思路与 prompt 缓存有相似之处：语言里很多片段会重复出现。

它的核心算法 _find_longest_matched_ngram_and_propose_tokens 是 KMP（Knuth-Morris-Pratt）字符串匹配算法的一个变体：

1
def _find_longest_matched_ngram_and_propose_tokens(
2
    origin_tokens,   # 已生成的 token 序列
3
    min_ngram, max_ngram,   # 匹配的 n-gram 长度上下界，默认都是 5
4
    max_model_len, k,       # k = 要提议的 token 数
5
) -> np.ndarray:
6
    # 1. 翻转 token 序列：tokens = origin_tokens[::-1]
7
    #    这样做的原因是我们要从"当前结尾"往前找共同后缀
8
    # 2. 构造 lps（Longest Proper Prefix-Suffix）数组，长度上限 max_ngram
9
    #    ——因为匹配段被 max_ngram 限制，只需要存前缀长度不超过 max_ngram 的 lps
10
    # 3. 在翻转后的序列里，找到长度落在 [min_ngram, max_ngram] 区间内的最长匹配后缀
11
    # 4. 返回"匹配段之后"接下来 k 个 token 作为草稿

为什么用翻转 + KMP？因为我们要找的是「当前序列的尾部后缀，在更早的历史里有没有出现过相同的」。把序列倒过来，就成了「当前序列头部的前缀」匹配问题，正好能套 KMP 的 lps 数组。匹配长度被 max_ngram 截断是为了省内存——lps 数组只需 max_ngram 长。

它从 SpeculativeConfig 读两个参数：

ngram 草稿有个明显的好处：为零成本，草稿前向的开销几乎可以忽略（纯 CPU 字符串查找，还能用 numba JIT 并行批量做）。但它只对「大量重复文本」（比如代码、格式化输出、对话模板）有效，对自由文本效果很差，所以实际生产用得多的还是模型类草稿。

DraftModelProposer：独立的草稿 LLM#

DraftModelProposer（vllm/v1/spec_decode/draft_model.py）对应用一个小得多的、独立的 LLM 当草稿模型的场景。这是最直观的投机解码形态：小模型算得快，猜 $K$ 个 token，大模型验证。

它有两个值得一提的点：

第一，草稿模型的 token 不共享目标模型的注意力/词表。 小模型和大模型往往是两个不同的 tokenizer。于是 DraftModelProposer 支持异构词表（heterogeneous vocab）：当 use_heterogeneous_vocab=True 时，会构建一个 VocabMapping（vllm/v1/spec_decode/vocab_mapping.py），在两个 tokenizer 的 token id 之间做双向映射。vLLM 用了一个叫 TLI（Token-Level Intersection） 的思路：只让草稿 logits 在「两个词表都共享的 token」上计算，避免把「草稿词表里的 token id」直接当成「目标词表里的 token id」用导致错位。

第二，张量并行度有硬约束。 代码里有 _raise_if_draft_tp_mismatch()：

1
def _raise_if_draft_tp_mismatch(self, draft_tp_size: int) -> None:
2
    target_tp_size = self.vllm_config.parallel_config.tensor_parallel_size
3
    if draft_tp_size != target_tp_size:
4
        raise ValueError(
5
            "draft_tensor_parallel_size ... 必须与目标模型的 "
6
            "tensor_parallel_size 一致 ..."
7
        )

原因与 torch.compile 缓存、以及验证阶段要把草稿 logits 和 target logits 对齐有关——TP 不一致会让通信分片对不上。这是我们后面「局限」一节的重要一条。

EagleProposer：依赖目标模型特征的草稿#

EagleProposer（vllm/v1/spec_decode/eagle.py）是社区里最常用的草稿方法，对应 EAGLE/EAGLE3（以及 DeepSeek 系的 MTP）。它和前面的最大区别是 pass_hidden_states_to_model=True：

1
class EagleProposer(SpecDecodeBaseProposer):
2
    def __init__(self, vllm_config, device, runner=None):
3
        super().__init__(
4
            vllm_config, device,
5
            pass_hidden_states_to_model=True,
6
            runner=runner,
7
        )

EAGLE 的核心洞察是：与其用一个小模型去猜「下一个 token 是什么」，不如把目标模型已经算好的 last hidden state 也喂给草稿层，让草稿层去学习「给定特征，下一步最可能是谁」。这样草稿模型只需做一次额外的低维映射，质量远高于从零猜。所以 EAGLE 需要一个「特征注入」的通道——target_hidden_states 必须从目标模型最后一层提取出来传给草稿模型。这一点直接体现在 propose() 里：

1
if self.method in ("eagle3", "dflash"):
2
    model = ... # 解包 BreakableCUDAGraphWrapper
3
    assert isinstance(model, (Eagle3LlamaForCausalLM, ...))
4
    target_hidden_states = self.model.combine_hidden_states(target_hidden_states)
5
    assert target_hidden_states.shape[-1] == self.hidden_size

EAGLE3 更进一步，支持 use_aux_hidden_state：不只从最后一层，还能从中间若干层提取特征（eagle_aux_hidden_state_layer_ids 配置）。站内的 EAGLE 系列文章 把特征级草稿、树验证、EAGLE-2/3 演进讲得很透了，这里我们只关注 vLLM 怎么把它装进去：核心就是「提取目标 hidden state → 拼进草稿输入 → 草稿模型前向」。

这张图来自 EAGLE 原始论文，展示了「草稿层（draft layer）」如何叠在目标模型最后一层隐藏状态之上做低维映射来预测下一 token：训练时（上图）草稿层用目标模型的特征做监督学习，解码时（下图）把上一轮已接受 token 的特征继续喂给草稿层，让它多步生成草稿序列。vLLM 的 EagleProposer 正是把「提取最后一层 hidden state」这一环接进了主循环。

MTP（Multi-Token Prediction）在 vLLM 里也走 EAGLE 这条路：DeepSeek-V3/V4 等模型的 model_type 落在 MTPModelTypes 这个 Literal 里（deepseek_mtp、qwen3_5_mtp、minimax_m3_mtp、glm4_moe_mtp、pangu_ultra_moe_mtp 等 20 多种），SpeculativeConfig.__post_init__ 会自动把 method 识别为 "mtp"，然后用 EAGLE 风格的 EagleProposer（实际上是同一套 Proposer 逻辑）加载这些预训练好的 MTP 头。

MedusaProposer：并行的多个头部#

MedusaProposer（vllm/v1/spec_decode/medusa.py）对应 Medusa 那类方法：不用额外的 LLM，而是在目标模型最后一层 hidden state 上挂多个小 MLP 头，每个头负责预测「往后第 $i+1$ 个 token」，所以一次前向就能给出全部 $K$ 个草稿 token：

1
class MedusaProposer:
2
    def propose(self, num_speculative_tokens, target_hidden_states, ...):
3
        blocks = self.model(target_hidden_states)         # 多个 MLP head 并行
4
        logits = self.model.compute_logits(blocks)
5
        draft_tokens = torch.stack(
6
            [logit.argmax(dim=-1) for logit in logits], dim=1
7
        )
8
        return draft_tokens   # [batch_size, num_heads]

它有个限制：如果目标是 MoE 模型且开启了 EPLB（Expert Parallel Load Balancing），Medusa 不被支持（medusa.py 里有对应的断言）。这属于「草稿方法 × 调度特性」的组合限制。

方法自动检测：你只需给个模型名#

用户层面，vLLM 让这一切尽量「零配置」。SpeculativeConfig.__post_init__（vllm/config/speculative.py）会根据草稿模型的名字 / HF config 的 model_type自动推断 method，判断顺序大致是：

1
模型名/type 含 "eagle3"            → method = "eagle3"
2
模型名含 "dflash" / "dspark"       → method = "dflash" / "dspark"
3
hf_config.model_type == "medusa"   → method = "medusa"
4
model_type ∈ MTPModelTypes         → method = "mtp"   （deepseek_mtp, qwen3_5_mtp ...）
5
model_type == "mlp_speculator"     → method = "mlp_speculator"
6
其他（例如单独给一个小模型路径）  → method = "draft_model"（或 "ngram"，当指定 `--speculative-method ngram` 时）

所以典型用法是：vllm serve <target_model> --draft-model <draft_model_or_eagle_ckpt> --num-speculative-tokens 5，框架自行判断该用 EAGLE 还是独立草稿模型。ngram 则是 --speculative-method ngram --num-speculative-tokens 5，不需要 --draft-model。

我们把「草稿从哪来」这五种方式总结成下表：

Batch Expansion：从「序列」到「并行验证的批」#

草稿模型提议的是每个请求一条长度为 $K$ 的候选序列。要让目标模型一次性验证这 $K$ 个位置，vLLM 必须把这些候选 token「摊平」成一个更大的 token 批，让目标模型对所有请求 × 所有草稿位置做一次前向。

关键方法在 llm_base_proposer.py：

• set_inputs_first_pass(...)：构造第一次草稿前向的输入。对 needs_extra_input_slots=False 的方法（主要是 EAGLE，因为它不额外 pad，直接 shift input_ids + 插入 next token），输入相对规整；对 needs_extra_input_slots=True 的方法（draft_model、并行 drafting），需要调用 Triton kernel _copy_and_expand_eagle_inputs() 做输入展开，并给每个 request 补 extra_slots_per_request 个 padding slot。
• prepare_inputs_padded(...) / prepare_next_token_ids_padded(...)：把「下一步的 token id」预展开成 padding 的形状。

所谓 padded drafter batch，是 V1 为了驯服 CUDA Graph 做的选择：把每个 request 的草稿输入都 pad 到相同长度，这样每次前向的输入形状固定，能被 CUDA Graph 捕获。代价是多算一些 padding slot——在 compute_new_slot_mapping 里，被拒绝的 token 位置会被置为 PADDING_SLOT_ID（-1），表示「这个 KV slot 不用再读了」。

展开时对 padding slot 的处理很关键：is_rejected_token_mask 和 is_masked_token_mask 两个 mask 跟踪哪些位置是真实草稿、哪些是 padding。被 mask 掉的位置不会贡献最终 logits（我们后面会看到 SpecDecodeMetadata 专门记录哪些 logits 是「要采样的 target 位置」）。

SpecDecodeMetadata：验证阶段的「索引地图」#

验证阶段需要的所有对齐信息，都打包在一个 SpecDecodeMetadata dataclass（vllm/v1/spec_decode/metadata.py）里：

1
@dataclass
2
class SpecDecodeMetadata:
3
    draft_token_ids: torch.Tensor       # [num_tokens] 所有摊平的 draft token
4
    num_draft_tokens: list[int]         # [batch] 每个 request 的草稿 token 数（K_i）
5
    cu_num_draft_tokens: torch.Tensor   # [batch] 草稿数量的前缀和
6
    cu_num_sampled_tokens: torch.Tensor # [batch] 前缀和（含 bonus 数）
7
    target_logits_indices: torch.Tensor # [num_tokens] target logits 的取址索引
8
    bonus_logits_indices: torch.Tensor  # [batch] bonus token 的 logits 索引
9
    logits_indices: torch.Tensor        # [num_tokens + batch] 合并索引
10
    max_spec_len: int                   # max(num_draft_tokens)

cu_num_draft_tokens 这个前缀和数组是整段的枢纽：它像一个「分组指针」，让我们能在摊平的 token 串 [req0_tok0, req0_tok1, ..., req1_tok0, ...] 里，高效地知道每个 request 的草稿 token 落在哪个区间。V0 时代用 Trie 树来组织树形注意力的父子关系；V1 的 padded drafter batch 则直接用这个前缀和来索引，不再需要 Trie。

上面这张 trace 图同样来自 DeepMind 那篇原始论文：它以一条真实草稿序列为例，从左到右画出「草稿模型提议 → 目标模型验证 → 逐位置 accept / reject → 拒绝处 recovered」的完整时间线。最上面是草稿 token，中间是目标模型对每个位置的检查，下面标出最终保留下来的接受序列——能直观对应后面要讲的拒绝采样器逻辑。

拒绝采样：RejectionSampler 的两个分支#

验证真正发生的地方是 RejectionSampler（vllm/v1/sample/rejection_sampler.py）。它是个 nn.Module，forward 大体分四步：

1. Bonus token：从 target logits 的「序列最后位置」再采一个额外 token（bonus_logits_indices 对应）。为什么需要 bonus？ 因为我们一共提议了 $K$ 个草稿，理论上可能全部被接受——那样这轮就只「消费」了 $K$ 个草稿 token 而没有真正产出目标模型的 token。为了补齐这一格（让每轮至少前进一步、避免死锁），在全部接受时追加一个由目标模型亲自采样的 token。这就是 bonus 的由来。

2. Target logits 提取与预处理：用 target_logits_indices 从目标模型整块 logits 里把「每两个相邻草稿之间该验证的那个位置」的 logits 抠出来，然后施加采样约束（temperature、top-k、top-p），再应用 logits processor（penalty 等）。注意这里要处理 num_draft_tokens 的重复展开——比如一个 request 有 $K_i$ 个草稿，那么它的温度/penalty 参数也要重复 $K_i$ 次（expand_batch_to_tokens）。

3. 拒绝采样主体（rejection_sample）分两套 kernel：

   • greedy 路径（rejection_greedy_sample_kernel）：对所有「贪婪采样」的请求，直接比较 draft_token_id == target_argmax_id。相同就接受，不同就拒绝并记下那个位置。不需要算概率。
   • 随机路径（rejection_random_sample_kernel）：对做随机采样的请求，按论文的无偏准则，逐位置比较 target_probs / draft_probs >= uniform，被拒绝的位置用一个「重采样 token（recovered token）」顶上。

4. parse_output：把 output_token_ids（初始填 PLACEHOLDER_TOKEN_ID 即 -1）里真正有效的 token 过滤出来，转成 list[list[int]] 返回。

我读源码时截取到的主流程伪代码，正好能对应上面四步：

1
def rejection_sample(draft_token_ids, num_draft_tokens, max_spec_len,
2
                     cu_num_draft_tokens, draft_probs, target_logits,
3
                     bonus_token_ids, sampling_metadata, ...):
4
    # 输出 buffer: [batch, max_spec_len + 1]（+1 留给 bonus），用 -1 占位
5
    output_token_ids = torch.full((batch_size, max_spec_len + 1), PLACEHOLDER_TOKEN_ID)
6

7
    # A. 需要 uniform 随机数（synthetic 模式 / 非全 greedy 时）
8
    uniform_probs = generate_uniform_probs(...)
9

10
    # B. greedy 请求：直接比 argmax
11
    if not sampling_metadata.all_random:
12
        target_argmax = target_logits.argmax(dim=-1)
13
        _rejection_greedy_sample(output_token_ids, ...)
14

15
    # C. 随机请求：先算 target 概率分布，采样 recovered tokens
16
    target_probs = target_logits.softmax(dim=-1, dtype=torch.float32)
17
    recovered_token_ids = sample_recovered_tokens(...)
18

19
    # D. 随机拒绝采样 kernel，把 accepted/recovered/bonus 装配进输出
20
    _rejection_random_sample(output_token_ids, ...)

Greedy kernel：逐 token 比较#

greedy 版本对应 rejection_greedy_sample_kernel 这个 Triton kernel，对每个 request（一个 program），从 cu_num_draft_tokens 推出 start_idx / end_idx，然后在这段里逐 position 比较：

1
@triton.jit
2
def rejection_greedy_sample_kernel(...):
3
    req_idx = tl.program_id(0)
4
    start_idx, end_idx = ...        # 从 cu_num_draft_tokens[req_idx] 算
5
    num_draft_tokens = end_idx - start_idx
6
    rejected = False
7
    for pos in range(num_draft_tokens):
8
        if not rejected:
9
            draft_id  = tl.load(draft_token_ids_ptr + start_idx + pos)
10
            target_id = tl.load(target_argmax_ptr + start_idx + pos)
11
            rejected = (draft_id != target_id)      # 一看到不一致就停
12
            tl.store(output_token_ids_ptr + ..., target_id)
13
    # 若全部接受（没 rejected），追加 bonus token
14
    if not rejected:
15
        bonus = tl.load(bonus_token_ids_ptr + req_idx)
16
        tl.store(output_token_ids_ptr + ..., bonus)

注意那个 rejected 标志：一旦遇到第一个不一致的 token，后面的草稿全部作废（这是投机解码的正确性要求——如果第 3 个 token 被拒绝，第 4 个 token 本身也是基于错误的上下文猜的，不能再用）。但 greedy 路径里，被接受的那些位置输出的是 target_argmax（而不是草稿 token），因为 target argmax 才真正代表「目标模型的决策」。只有全部接受时才需要 bonus token 来「补一格」。

Random kernel：无偏接受 + 重采样#

随机路径对应 rejection_random_sample_kernel，它实现的就是我们背景小节里那两个公式：

• 接受：uniform < target_prob / draft_prob；
• 拒绝：从调整后的分布重采样一个 recovered_token_id。

sample_recovered_tokens 用到一个技巧：标准的多项式采样可以用「指数分布随机数求 argmax」来实现（Gumbel 思想的变体）。它给每个位置生成独立的指数随机数 $q$，对每个候选 token 算 adjusted_prob * (1/q)，取 argmax 就等价于「从 adjusted 分布采样」。adjusted prob 的分量是：

也就是把目标分布里「草稿 token $x$ 那一格」的质量整个剔除（因为那个 token 已经在这轮被判断为不接受了），再归一化。这保证了重采样不会又采回那个被拒的草稿 token，且最终分布无偏。

关于 synthetic / block 验证模式#

rejection_sample_method 除了默认的 standard，还有两档：

• synthetic：给每个位置预设一个「合成接受率」synthetic_acceptance_rates（必须单调非增），直接用 uniform < rate 接受草稿，不依赖真实的 target probs。它用于在还没有真实草稿模型时模拟评估（比如做实验/基准测算），不是生产用法。
• block：块验证（Block verification），按 Sun 等人的 Block Verification Accelerates Speculative Decoding，一次性联合验证整段草稿块而不是逐 token，可以省掉部分概率比较。

propose() 的多步循环#

SpecDecodeBaseProposer.propose() 的骨架（我在源码里核对过）大致是：

1
def propose(self, num_speculative_tokens, target_token_ids, target_positions,
2
            target_hidden_states, next_token_ids, ...):
3
    # —— Step 0：第一次前向 ——
4
    num_tokens, token_indices_to_sample, common_attn_metadata = (
5
        self.set_inputs_first_pass(target_token_ids=target_token_ids,
6
                                   next_token_ids=next_token_ids, ...))
7
    per_layer_attn_metadata = self.build_per_group_and_layer_attn_metadata(...)
8
    cudagraph_runtime_mode, ... = self._determine_batch_execution_and_padding(num_tokens)
9

10
    with set_forward_context(per_layer_attn_metadata, self.vllm_config,
11
                             cudagraph_runtime_mode=...):
12
        last_hidden_states, hidden_states = self.model(**model_kwargs)   # 草稿第 1 步
13

14
    # K==1 或并行 drafting：一次前向即可，直接采样后返回
15
    if self.num_speculative_tokens == 1 or self.parallel_drafting:
16
        draft_token_ids, draft_probs = self._sample_draft_tokens(
17
            last_hidden_states[token_indices_to_sample], sampling_metadata)
18
        return draft_token_ids.view(-1, K)
19

20
    # —— Step 1..K-1：逐步自回归 ——
21
    draft_token_ids_list = []
22
    for token_index in range(self.num_speculative_tokens - 1):
23
        positions = self._update_positions_dependent_metadata(...)   # 更新位置
24
        # 用 draft_index 感知重建注意力元数据（每个 step 位置不同）
25
        _, per_layer_attn_metadata = self.build_per_group_and_layer_attn_metadata(
26
            common_attn_metadata, draft_index=token_index + 1)
27
        # 把上一步的 input_ids / hidden_states 拷进 CUDA Graph buffer
28
        self.input_ids[:batch_size] = input_ids
29
        self.hidden_states[:batch_size] = hidden_states
30
        with set_forward_context(...):
31
            ret_hidden_states = self.model(**model_kwargs)          # 草稿第 N 步
32
        draft_token_ids, draft_probs = self._sample_draft_tokens(
33
            last_hidden_states[:batch_size], sampling_metadata)
34
        draft_token_ids_list.append(draft_token_ids)
35
    return torch.stack(draft_token_ids_list, dim=1)   # [batch, K]

几个要点：

1. draft_index：每一步的位置/attention 元数据都不同（第 $j$ 步的草稿 token 只 attend 到它前面的 token），所以多步循环里要按 draft_index 重建 attention metadata，而不是复用第 0 步的。
2. 中间状态存哪：V1 把多步所需的中间 hidden_states、input_ids、positions、两个 mask 都作为 SpecDecodeBaseProposer 的预分配 buffer 持久保存（__init__ 里按 max_num_tokens 分配），循环里就地更新，而不是每次新建张量——这是为了减少分配、贴合 CUDA Graph 的固定地址要求。
3. _sample_draft_tokens：每步从前一步的 hidden state 采样出下一个草稿 token，draft_sample_method 控制是 greedy（始终 argmax，速度最快）还是 probabilistic（从分布随机采样，能拿到完整的 draft_probs 供随机拒绝采样做概率比）。

CUDA Graph 在投机解码里怎么工作#

前面第七、九篇讲过 CUDAGraph 需要固定的 batch shape。草稿模型的多步前向天然和这个冲突：每一步的输入长度都在变。vLLM 的处理是靠两个动作：

• _determine_batch_execution_and_padding：在每次草稿前向之前，根据当前 num_tokens 算出该用哪个 CUDA Graph（cudagraph_runtime_mode 可以是 PIECEWISE 等），并把 batch pad 到图要求的形状。这就是 padded drafter batch 的用武之地。
• CudagraphDispatcher / initialize_cudagraph_keys / dummy_run：捕获阶段跑一次空转（dummy_run）来把形状固定下来；CudagraphDispatcher 根据 batch size 分派到正确的图。

有意思的是，多步循环里对 is_graph_capturing or self.parallel_drafting 走的是「只 forward 一次」的简化路径——因为捕获阶段只需要一步的形状；正常执行时才循环 K-1 次。这避免了在捕获阶段重复展开。

KV Cache 与位置在每一步的更新#

多步草稿的 KV Cache 管理是另一个重点。每生成一个草稿 token，都会给它分配一个 KV slot；但如果这个 token 后面被拒绝，它的 KV 就浪费了。vLLM 用 _update_positions_dependent_metadata 和 compute_new_slot_mapping 这两个工具函数来动态更新 slot mapping：被拒绝的草稿位置 → PADDING_SLOT_ID（-1），不会再分配/读取。这样草稿 KV 是一个「可回滚的临时区」。

最小启用法#

用独立草稿模型：

1
vllm serve meta-llama/Meta-Llama-3-8B-Instruct \
2
    --draft-model JackFram/llama-68m \
3
    --num-speculative-tokens 5

用 EAGLE 草稿（需要一个带 EAGLE 权重/映射的 checkpoint）：

1
vllm serve meta-llama/Meta-Llama-3-8B-Instruct \
2
    --draft-model Yuhuili/EAGLE-LLaMA3-Instruct-8B \
3
    --num-speculative-tokens 5

用 ngram（不加载草稿模型）：

1
vllm serve meta-llama/Meta-Llama-3-8B-Instruct \
2
    --speculative-method ngram \
3
    --num-speculative-tokens 5 \
4
    --ngram-prompt-lookup-max 5

关键参数（对应 SpeculativeConfig 字段）：

Dynamic SD：按 batch size 调 K#

SpeculativeConfig 里有个 num_speculative_tokens_per_batch_size（格式 [(range_start, range_end, K), ...]），对应 dynamic speculative decoding：batch 很大时，模型利用率已经很高，再投机收益不大，可以调小 $K$ 甚至到 0；batch 小时，decode 是内存瓶颈，投机收益大，调大 $K$。vllm/v1/spec_decode/dynamic/utils.py 里的 build_dynamic_sd_schedule_lookup 会把这组区间编译成一个 O(1) 的 dense_schedule[batch_size] → K 查表。

配置校验与常见报错#

_verify_args() 做了不少约束检查，常见的：

• synthetic_acceptance_rates 与 synthetic_acceptance_length 互斥，且 rates 必须单调非增。
• use_heterogeneous_vocab=True 时，draft_sample_method 必须为 greedy，且只对 method="draft_model" 有效。
• draft_tensor_parallel_size 必须等于目标 tensor_parallel_size。
• DeepSeek-V4.1 的 checkpoint 不含经典 MTP 权重，必须用 dspark 而不是 mtp（这是模型侧的特殊情况）。