zero-to-sglang

简体中文

English

简体中文

English

外观

第 10 章 投机解码（代码走读） ​

普通生成每调用一次目标模型，只确认一个新 token。投机解码先提出多个候选，再由目标模型验证。本篇承接概念篇，按“起草 → 验证 → 提交”的顺序阅读代码，重点看位置对齐、首次拒绝和停止条件。

前面章节引用的 mini-sglang 9a91cfa 没有投机解码模块，因此本篇定位真实源码时使用 SGLang 的固定提交 1291f317。代码块分别标明源码简化或教学算法片段，需要结合所述接口阅读，不是可独立启动的推理程序。

1 本章学习目标 ​

完成本章学习后，你将能够：

  1. 定位草稿模型、验证输入和执行流程的源码入口。
  2. 找准每个草稿对应的预测行，处理首次拒绝与额外 token。
  3. 解释随机路径中的接受概率和拒绝校正，以及缓存需要恢复的状态。

2 定位源码 ​

以下路径均相对于 SGLang 的 python/sglang/srt/speculative/：

文件与入口阅读重点
standalone_worker_v2.py：StandaloneDraftWorker、StandaloneWorkerV2独立草稿模型及词表兼容性检查
eagle_worker_v2.py：EAGLEWorkerV2.forward_batch_generation起草、目标验证与草稿状态更新的组织顺序
eagle_info.py：EagleVerifyInput草稿 token、位置、检索索引与验证元数据

独立草稿复用部分 EAGLE 执行基础设施，但二者的模型机制并不相同。下面先用一条线性草稿理解算法，再回到这些入口；树状候选与 GPU 内核的细节不在本篇展开。

3 先看起草与验证怎样衔接 ​

3.1 源码中的调用顺序 ​

在 forward_batch_generation 的普通起草路径中，关键顺序可以简化为：

python
# eagle_worker_v2.py: EAGLEWorkerV2.forward_batch_generation (simplified)
verify_input = self.draft_worker.draft(batch)
batch.spec_info = verify_input
batch_output = self.verify(batch)
self.draft_worker._draft_extend_for_decode(batch, batch_output)

这里省略了并行上下文、发布回调和特殊分支。draft 准备候选与验证元数据，verify 调用目标模型并确定接受结果，最后更新草稿侧状态，为下一轮准备输入。沿 verify 还能继续找到 run_eagle_verify。

3.2 线性草稿逐步依赖前面的候选 ​

下面改用教学接口：tokens 是已确认的非空 token ID 列表，draft(prefix) 返回每个输入位置之后的 logits，形状为 [len(prefix),V]。target 与 draft 必须使用相同的 token ID 含义。draft_steps 为正整数，remaining 是剩余输出额度，eos_id 是结束标记的 ID（不使用时为 None）。仅在还有输出额度时进入下一轮。贪心起草片段为：

python
# Teaching fragment: greedy linear drafting
proposals = []
for _ in range(min(draft_steps, remaining)):
    token = int(draft(tokens + proposals)[-1].argmax())
    proposals.append(token)
    if token == eos_id:
        break

候选暂存在 proposals 中，不能直接写进已确认输出。草稿提出 EOS 时停止继续起草，但这个 EOS 仍要由目标模型验证。

3.3 一次目标前向得到 k+1 行 ​

target(prefix) 使用同样的全位置 logits 接口。令已确认前缀长为 L，草稿长为 k，则：

从 0 开始的 target 行号用途
L-1验证第一个草稿
L验证第二个草稿
L+k-2验证最后一个草稿
L+k-1全接受后选择额外 token（bonus token）

教学片段按下面的区间取出这些行：

python
# Teaching fragment: full-prefix verification without a persistent KV cache
verification = target(tokens + proposals)[len(tokens) - 1:]

第 i 行预测第 i+1 个 token，所以起点是 L-1。因果掩码保证前面的预测行看不到后面的草稿。真实引擎可以复用已缓存的前缀；这里使用全前缀接口是为了让位置对应清楚可见。

4 贪心验证：在首次拒绝处停止 ​

4.1 提交接受前缀和纠正 token ​

下面的教学函数只负责一轮提交：verification 是上一节的 [k+1,V] 张量，remaining 为非负整数，返回值是本轮确认的 token 列表。

python
# Teaching fragment: greedy commitment for one verification round
def greedy_commit(proposals, verification, remaining, eos_id=None):
    if remaining == 0:
        return []
    output = []
    for i, proposed in enumerate(proposals):
        token = int(verification[i].argmax())
        output.append(token)
        if token == eos_id or len(output) == remaining:
            return output
        if token != proposed:
            return output
    output.append(int(verification[len(proposals)].argmax()))
    return output

若目标选择与草稿相同，就继续检查下一行；首次不同则提交目标选择，立即返回。全接受且仍有输出额度时，最后一行再提供一个 bonus token。调用方把返回列表追加到已确认前缀；若最后一个 token 是 EOS，或额度已用完，就结束生成，包括 bonus 为 EOS 的情况。

4.2 手算一次拒绝 ​

假设草稿提出 [1,1,1]，目标在前两行分别选择 1、2。第一个 1 被接受，第二个 1 被纠正为 2，本轮只提交 [1,2]。剩余草稿的预测以错误的前缀为条件，必须丢弃，不能继续使用后面的验证行。

贪心正确性的关键是：每次提交的 token 都是目标在当前已确认前缀上的选择。固定相同的数值计算和并列分数处理方式时，输出应与目标单独逐步生成一致。

5 随机验证与缓存边界 ​

5.1 接受概率与正残差 ​

随机路径中，候选 y 从草稿分布 q 抽出，以 min(1, p[y] / q[y]) 的概率接受；拒绝后从正残差归一化得到的分布抽取纠正 token。以下教学片段假定 p、q 是同一条件前缀下、相同词表上的一维概率张量，且 q[y] > 0：

python
import torch

acceptance = min(1.0, float(p[y] / q[y]))
if torch.rand(()).item() < acceptance:
    token = y
else:
    residual = (p - q).clamp_min(0)
    token = int(torch.multinomial(residual / residual.sum(), 1))

这里的 p、q 必须是实际使用的采样分布；若应用温度、Top-k 或 Top-p，就保存处理后的概率。理论上发生拒绝时正残差总量大于零；工程实现还需检查有限精度下的数值有效性。

例如 p=[0.4,0.6]、q=[0.8,0.2]：提出 0 时接受概率为 0.5，提出 1 时总是接受，拒绝后只能补 1。最终输出 0 的概率是 0.8×0.5=0.4，输出 1 的概率是 0.2+0.8×0.5=0.6。拒绝后直接从 p 重抽会破坏这个结果。

接受分支贡献 min(p(v),q(v))，拒绝分支补齐剩余的正差；每一步使用正确的条件前缀，才能保持整个序列的目标分布。首次拒绝后仍要停止本轮，全接受后的 bonus 则从目标分布抽取。随机路径保证的是分布一致，不是相同 seed 下逐 token 相同。

5.2 回到验证输入与 KV Cache ​

打开 EagleVerifyInput，重点看 draft_token、positions、custom_mask 和 retrieve_*：真实验证需要记录候选的位置与可见关系，不能仅拼接 token IDs。启用拒绝采样时，draft_probs 保存逐步草稿分布，供验证使用。

本篇教学片段没有持久化 KV。真实引擎可能已为未接受草稿写入缓存，首次拒绝后需要恢复到已确认路径。纠正 token 或 bonus token 只是刚被选出，通常尚未完成自己的前向，不能把被拒绝候选的 KV 当成它们的缓存。沿执行入口继续阅读时，应同时追踪接受长度、缓存索引和草稿侧状态更新。

目标调用次数减少不等于同等比例的加速。实际收益还取决于草稿耗时、接受长度、验证与缓存管理开销，需要在真实模型和硬件上测量。

6 总结与测试题 ​

6.1 课程总结 ​

投机解码把起草与确认分开。验证行必须对齐候选之前的前缀，提交只保留连续接受的部分及必要的纠正 token。贪心路径比较目标选择，随机路径使用接受概率与正残差；EOS、长度上限和缓存恢复都必须围绕已确认输出处理。

6.2 测试题 ​

  1. 前缀长 5、草稿长 3，验证三个草稿和选择 bonus 分别读取哪些行？
  2. 为什么首次拒绝后不能继续使用后面的验证行？
  3. 在第 5.1 节的例子中，拒绝后直接从 p 抽样，会得到怎样的输出概率？
  4. 草稿提出 EOS 但目标拒绝了它，此时应提交什么，是否结束请求？
  5. 为什么词表大小相同仍不足以说明 target 与 draft 兼容？
  6. 为什么纠正 token 不能直接复用被拒绝候选的 KV？

参考资料 ​