Rust in Everything

投机解码

主线位置:一个请求的一生里「前向 → 采样」这一段,解码阶段每一步都会经过。草稿在采样之后提出,下一步和目标模型的前向一起验证。两个引擎这一段都是 Python 编排加 CUDA / Triton kernel,没有 Rust;页末拿纯 Rust 的 mistral.rs 作对照。
版本:vLLM v0.32.0、SGLang v0.5.22

结论先说:

  • 投机解码不改模型,只改「一步出几个 token」。便宜的草稿先猜 k 个,目标模型一次前向算 k+1 个位置的分布,从左往右接受;第一个不接受的位置用目标分布补一个 token,全接受就多送一个。一轮最少出 1 个 token,最多出 k+1 个。
  • 两边默认都不开。vLLM 用 --speculative-config 传一个 JSON;SGLang 用 --speculative-algorithm 加一组 --speculative-* 参数。方法上两边都有 n-gram、独立草稿模型、EAGLE / EAGLE3、模型自带的 MTP 头,各自还有几种更新的方法。
  • 接受规则保证输出分布和不开投机时一样。贪心时就是逐 token 比 argmax;采样时是拒绝采样:以 min(1, p/q) 接受草稿,拒绝时从 max(0, p−q) 归一化后重采。两边的默认实现都不用草稿的概率:vLLM 默认把草稿当 one-hot(draft_sample_method="greedy"),SGLang 默认走只看目标概率的 TreeSpeculativeSamplingTargetOnly。这两种是拒绝采样的特例,仍然无损。
  • SGLang 有两个放宽接受的阈值,默认 1.0;调小会提高接受率,但输出分布就变了。
  • 草稿 token 要先占 KV 槽。vLLM 按 num_computed_tokens 回退,块不还;SGLang 每步为每个请求预留 2 × max(步数 × topk, 草稿 token 数) 个槽位,结束时把多占的还回去。SGLang 开投机后,没设 --max-running-requests 时会被设成 48。
  • 接受率、加速比、开关前后的输出一致性都待实测(C3-5)。
flowchart LR
  prev["上一步采样出的 token"] --> draft["草稿提出 k 个 token<br/>n-gram 查表 / 草稿模型 / EAGLE 头"]
  draft --> sched["调度器把 1 + k 个 token<br/>排进下一步,分配 KV 槽"]
  sched --> fwd["目标模型一次前向<br/>得到 k+1 个位置的分布"]
  fwd --> verify{"从左往右验证"}
  verify -->|"第 i 个被拒"| fix["用目标分布补 1 个 token<br/>后面的草稿作废"]
  verify -->|"k 个全接受"| bonus["再从第 k+1 个位置采 1 个"]
  fix --> roll["回退:已算 token 数减去被拒数"]
  bonus --> out["本轮输出 1 到 k+1 个 token"]
  roll --> out
  out --> prev

一轮出几个 token

解码访存受限(见 为什么服务大模型难):一步主要花在把权重读一遍上。验证 k+1 个位置和算 1 个位置读的权重一样多,所以只要草稿够便宜、接受得够多,一步就能多出几个 token。

如果每个草稿 token 被接受的概率都是 α、彼此独立,一轮期望输出的 token 数是 (1 − α^(k+1)) / (1 − α)。α = 0.7、k = 3 时约 2.53 个(估算,只是这个公式的算术;没算草稿本身的开销,也没算验证前向比普通解码前向慢多少)。真实的接受率随模型、草稿、提示和温度变,待实测(C3-5)。

两边支持哪些方法

方法vLLM v0.32.0SGLang v0.5.22草稿从哪来
n-gramngram、ngram_gpuNGRAM不用模型:在已有 token 里找和结尾相同的片段,取它后面的 token
独立草稿模型draft_model(不写 method 时的默认推断)STANDALONE一个同词表的小模型自回归跑 k 步
EAGLE 系eagle、eagle3EAGLE、EAGLE3一层小网络,输入目标模型的隐藏状态,自回归跑 k 步
模型自带 MTP 头mtp(旧的 deepseek_mtp 等名字都改写成 mtp)NEXTN(别名,解析成 EAGLE)checkpoint 里自带的预测层,按 EAGLE 的方式跑
其他medusa、suffix、dflash、dspark、custom_class 等DFLASH、DSPARK、UNO,以及插件注册的算法各自论文的做法,本页不展开

出处:vLLM 的 SpeculativeMethod 和 MTPModelTypes,MTP 别名在 speculative.py#L1143 改写;SGLang 的 SpeculativeAlgorithm 和别名表 _RESERVED_ALIASES。

vLLM 这边还要看 model runner。v0.32.0 在 CUDA 上默认用 V2 model runner(VLLM_USE_V2_MODEL_RUNNER 不设时,有 Triton 且没有不支持的特性就选 V2,vllm.py#L703)。V2 不支持 suffix、medusa、mlp_speculator、custom_class,遇到它们就退回 V1 runner(vllm.py#L2929)。本页讲的 vLLM 代码都是 V2 的 vllm/v1/worker/gpu/spec_decode/;V1 runner 有自己的一套提议器和 vllm/v1/sample/rejection_sampler.py,接受规则相同。

V2 按方法建「speculator」:

vllm/v1/worker/gpu/spec_decode/__init__.py

def init_speculator(
    vllm_config: VllmConfig,
    device: torch.device,
    req_states: "RequestState",
):
    # ...
    elif speculative_config.method == "mtp":
        # ...
        return MTPSpeculator(vllm_config, device)
    elif speculative_config.use_eagle():
        # ...
        return EagleSpeculator(vllm_config, device)
    elif speculative_config.uses_draft_model():
        # ...
        return StandaloneARSpeculator(vllm_config, device)
    elif speculative_config.use_ngram():
        # ...
        return NgramGPUSpeculator(vllm_config, device, req_states)

SGLang 按算法返回 worker 类:EAGLE / EAGLE3 / NEXTN 用 EAGLEWorkerV2,STANDALONE 用 StandaloneWorkerV2,NGRAM 用 NGRAMWorker(spec_info.py#L302)。

怎么配置

vLLM:一个 JSON

--speculative-config(简写 -sc)接一个 JSON,字段就是 SpeculativeConfig 的字段(arg_utils.py#L1828):

# EAGLE3:草稿头单独一个 checkpoint
vllm serve <目标模型> \
  --speculative-config '{"method": "eagle3", "model": "<EAGLE3 头>", "num_speculative_tokens": 3}'

# n-gram:不需要草稿权重
vllm serve <目标模型> \
  --speculative-config '{"method": "ngram_gpu", "num_speculative_tokens": 4}'
字段默认说明
methodNone不写时:model 是 ngram 就当 n-gram,否则当 draft_model(speculative.py#L1139)
num_speculative_tokensNone(#L393)草稿配置里有就用它的,否则必须给
prompt_lookup_min / max两个都不给时都设成 5(#L1199)n-gram 的匹配长度
draft_sample_method"greedy"(#L600)草稿取 argmax;"probabilistic" 才按草稿分布采样并缓存草稿 logits
rejection_sample_method"standard"(#L529)另有 block(整块验证)和 synthetic(按设定的接受率随机接受,用于压测,不保证分布)

还有两处间接影响:

  • 模型目录本身是 speculators 格式的草稿 checkpoint 时,vLLM 会自动改写成「目标模型 + 投机配置」(arg_utils.py#L2230)。
  • 异步调度默认开,但方法不在 EAGLE 系、ngram_gpu、draft_model、dspark、dflash 里时会被关掉(vllm.py#L1600)。判断按方法名:写 ngram 关,写 ngram_gpu 不关;而在 V2 runner 里两者走同一个 GPU 实现(use_ngram() 同时认两个名字,speculative.py#L1979)。

SGLang:一组参数,部分启动时推出来

python -m sglang.launch_server --model-path <目标模型> \
  --speculative-algorithm EAGLE3 --speculative-draft-model-path <EAGLE3 头>
参数默认说明
--speculative-algorithmNone(spec.py#L31)不设就不开
--speculative-num-steps / --speculative-eagle-topk / --speculative-num-draft-tokensNone(spec.py#L58)EAGLE 系三个都不给时按模型架构挑:LlamaForCausalLM 是 (5, 4, 8),其余多数是 (3, 1, 4)(speculative_hook.py#L1340);topk 为 1 时草稿 token 数强制等于步数 + 1(#L1170)
--speculative-accept-threshold-single / -acc1.0(spec.py#L109)放宽接受的两个阈值,见下文
--speculative-use-rejection-samplingFalse(spec.py#L117)打开后用草稿概率做标准拒绝采样;只支持 EAGLE / EAGLE3 / NEXTN 且 topk = 1(speculative_hook.py#L1124);ROCm 上会自动打开
NGRAM 专用草稿 token 数 12(#L1228),topk 取 --speculative-ngram-max-bfs-breadth(默认 10,spec.py#L259)n-gram 的树宽

(5, 4, 8) 的意思是:草稿跑 5 步,每步留 4 个候选,最后从这棵树里挑 8 个 token 送去验证。topk > 1 时验证的是一棵树,不是一条链。

开投机后的副作用:EAGLE 系和 NGRAM 在没设 --max-running-requests 时把它设成 48,并打警告(speculative_hook.py#L999、#L1208);混合块(--enable-mixed-chunk)对不支持它的算法关掉。并发上限从按显存推算变成 48,对 连续批处理 的影响比投机本身还直接。

草稿怎么提出

n-gram:不用模型

vLLM 的 GPU 版对每个请求单独查:拿已有 token 序列的结尾,在同一个请求的历史里找最长、最靠右的相同片段,取它后面的 k 个 token。kernel 把「长度」和「位置」打包成一个数,一次取最大值就同时满足两条:

vllm/v1/worker/gpu/spec_decode/ngram/speculator.py

        # Pack (n, pos) so a single max yields longest-n, rightmost-pos.
        cand = n_iter * Lp1 + pos + 1
        best_score = tl.where(match, cand, best_score)

SGLang 的 NGRAM 维护一棵跨请求的 trie(C++ 实现,容量默认 1000 万,spec.py#L271),每步把各请求新出的 token 放进去,再按请求结尾做 BFS,取出一棵候选树(ngram_worker.py#L299)。区别在于:vLLM 只能抄本请求自己说过的话,SGLang 能抄别的请求说过的话。

n-gram 的草稿没有概率分布,验证时只能用「只看目标概率」的规则。

EAGLE 与 MTP:读目标模型的隐藏状态

EAGLE 头和 MTP 层的输入是目标模型最后(EAGLE3 是几层拼接)的隐藏状态加上已采样的 token。所以它们必须在目标模型前向之后跑,每轮的顺序是:目标前向 → 验证采样 → 草稿第一步(复用目标这一步的隐藏状态)→ 草稿再跑 k−1 步。

V2 里第一步草稿把输入长度对齐到目标这一步的查询长度,被拒的位置也留着,这样 CPU 不用等 GPU 告诉它拒了几个:

vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py

        # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
        # number of rejected tokens, we maintain the size of input_ids and
        # hidden_states the same as the target model's. This means, we pad each
        # request's query length to include any rejected positions. By doing so,
        # we can also reuse the attention metadata (e.g., query_start_loc,
        # seq_lens) of the target model.

草稿每步怎么取 token:vLLM 默认直接 argmax,只有 draft_sample_method="probabilistic" 才按温度采样并把 logits 存下来给验证用(speculator.py#L439,argmax 在 #L473)。SGLang 每步取 topk 个候选组成树;topk = 1 时退化成一条链。

独立草稿模型

一个完整的小模型,有自己的 KV cache,自回归跑 k 步。它不依赖目标的隐藏状态,但要和目标共用词表(vLLM 有 use_heterogeneous_vocab 处理不同词表,V2 不支持,会退回 V1)。

目标模型怎么一次验证

调度器把「上一步采样出的 token + k 个草稿」当成这个请求本步的 1 + k 个新 token,和普通预填充块一样进前向。注意力按因果掩码算,第 i 个位置只看到前 i 个草稿,所以一次前向就得到 k+1 个条件分布:p(· | 前缀)、p(· | 前缀, d1)、……、p(· | 前缀, d1…dk)。SGLang 验证一棵树时,用自定义掩码让每个节点只看到自己的祖先。

V2 在采样处分岔:本步有草稿就走拒绝采样器,否则走普通采样器。

vllm/v1/worker/gpu/model_runner.py

            # Rejection sampling for spec decoding.
            assert self.rejection_sampler is not None
            assert self.speculator is not None
            sampler_output = self.rejection_sampler(
                logits,
                input_batch,
                # Draft logits are needed for probabilistic rejection sampling.
                self.speculator.draft_logits,
                invalid_drafts,
            )

draft_logits 只在 probabilistic 时不是 None。验证完,同一步接着让 speculator 提出下一轮草稿(model_runner.py#L2239)。

接受规则为什么不改变分布

记目标在某个位置的分布为 p(已经套过温度、top-k / top-p 等采样参数),草稿在同一位置的分布为 q,草稿 token 为 x。

贪心

温度为 0 时 p 是 argmax 上的单点分布,规则退化成:草稿等于目标 argmax 就接受,否则用目标 argmax 替换并停止。vLLM 在拒绝时直接把 argmax 写进结果,省掉重采:

vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py

            if is_greedy:
                # Greedy sampling. Accept IFF draft matches target argmax.
                # NOTE: Target argmax is stored directly so that resampling
                # can be skipped upon rejection.
                # ...
                    accepted = target_argmax == draft_sampled
                accepted &= is_valid_draft
                verifying = accepted

SGLang 在全贪心的批上走 verify_tree_greedy,比的也是 argmax(eagle_utils.py#L868)。数学上输出和不开投机一致;数值上验证前向的形状(每个请求 k+1 个位置)和普通解码不同,浮点误差可能让近似并列的 argmax 翻转,开关前后的输出一致性待实测(C3-5)。

采样:标准拒绝采样

规则:以 min(1, p(x)/q(x)) 接受 x;拒绝时从 r(y) = max(0, p(y) − q(y)) / Z 采一个 token,Z 是归一化常数。

推导只有两行。输出 y 的概率是「草稿就是 y 且被接受」加上「被拒后重采到 y」:

  • 第一项:q(y) · min(1, p(y)/q(y)) = min(p(y), q(y))。
  • 被拒的总概率:1 − Σ min(p, q) = Σ max(0, p − q) = Z,所以第二项是 Z · max(0, p(y) − q(y)) / Z = max(0, p(y) − q(y))。
  • 相加:min(p, q) + max(0, p − q) = p(y)。

顺带得到接受率 = Σ min(p, q) = 1 − TV(p, q),草稿越像目标接受得越多。多个位置时,对每个位置在「前面都已接受」的条件下用同一个论证,被拒后的草稿全部丢弃,全接受时从第 k+1 个位置的 p 再采一个,归纳下去整段输出的联合分布就是目标模型的自回归分布。

vLLM 的 kernel 在对数空间做这一比:

vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py

                    # Probability ratio test: p(x) > u * q(x)
                    # Equivalent log form: log_p(x) > log(u) + log_q(x)
                    accepted = target_logprob > tl.log(u) + draft_logprob

拒绝后的残差分布在 _resample_kernel 里算,同样取对数后交给 Gumbel 采样:

vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py

        #   r(x) = max(p(x) - q(x), 0)
        # Gumbel sampling needs logits, so we compute it in log space:
        #   log(r(x)) = log(max(exp(log_p(x)) - exp(log_q(x)), 0))
        # The more numerically stable form is:
        #   log(max(exp(a) - exp(b), 0)) = a + log(max(1 - exp(b - a), 0))
        ratio = tl.exp(draft_log_probs - target_log_probs)

SGLang 打开 --speculative-use-rejection-sampling 后走 chain_speculative_sampling_triton,判断式是 coin * q < p(reject_sampling.py#L149),同一个规则。

默认路径:草稿当 one-hot

vLLM 默认 draft_sample_method="greedy",草稿取 argmax,不存草稿 logits。这时把 q 当成 x 上的单点分布:q(x) = 1,接受概率变成 p(x);残差 max(0, p − δx) 就是把 x 的概率清零后的 p。kernel 里两处都写明了:

vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py

    else:
        # One-hot draft: q(token) = 1, log_q = 0.
        draft_log_prob = 0.0
        draft_lse = 0.0

vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py

        # One-hot draft. The residual is just the target distribution with
        # the rejected draft token probability zeroed out.
        # ...
        residual_logits = tl.where(
            block != rejected_draft_token,
            target_logits,
            float("-inf"),
        ).to(tl.float32)

这仍是上面的推导,只是 q 换成单点分布,所以照样无损。代价是接受率从 Σ min(p, q) 变成 p(x):草稿分布里的「把握」信息没用上。换来的是不用为每个请求、每个草稿位置存一份词表大小的 logits(probabilistic 要额外分配 max_num_reqs × k × vocab 的缓冲,speculator.py#L180)。

SGLang 默认:只看目标概率,可以验证一棵树

SGLang 不开拒绝采样时,把 draft_probs 置成全 0 传进 TreeSpeculativeSamplingTargetOnly(eagle_utils.py#L1003)。kernel 在树的每一层只抛一枚硬币,按兄弟节点的顺序累加目标概率:

python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh

      DType target_prob_single = target_probs[cur_prob_offset + draft_token_id];
      prob_acc += target_prob_single;

      const bool has_target_mass = target_prob_single > DType(0);
      if (has_target_mass && (coin < prob_acc / threshold_acc || target_prob_single >= threshold_single)) {
        // accept token
        // ...
        break;
      } else {
        // FIXME: leverage draft probs
        draft_probs[cur_prob_offset + draft_token_id] = target_probs[cur_prob_offset + draft_token_id];
        cur_index = retrive_next_sibling[bx * num_draft_tokens + cur_index];
      }

阈值都取默认 1.0 时,这就是按目标分布做逆 CDF 采样,只是把候选 token 排在前面:

  • 硬币 u 落在第 j 个兄弟的累计区间里,就接受第 j 个,概率正好是 p(c_j)。
  • 一个都没落中(概率 1 − Σ p(c_i)),被试过的兄弟在 draft_probs 里被填成目标概率,之后从 relu(p − draft_probs) 采样(sampling.cuh#L104),也就是把候选清零后归一化的 p。非候选 token y 的概率是 (1 − Σ p(c_i)) · p(y) / (1 − Σ p(c_i)) = p(y)。

所以默认路径同样无损,每层接受率是候选兄弟的目标概率之和。topk = 1(多数架构的默认)时只有一个候选,和 vLLM 的 one-hot 规则完全一样。n-gram 的草稿没有分布,也走这条路。

两个阈值是有意放宽的口子:coin < prob_acc / threshold_acc 在 threshold_acc < 1 时把接受区间放大,target_prob_single >= threshold_single 在 threshold_single < 1 时让目标概率够大的候选直接通过。两者都让候选 token 被选中的概率超过 p(c_j),输出分布随之偏离目标。默认都是 1.0,不会触发。

还有三处 SGLang 的默认行为和「严格等于不开投机」有出入,都写在代码注释里:

  • 验证时的重复惩罚是「放宽版」,按请求累计后统一加到所有草稿位置上(eagle_utils.py#L783)。
  • 在 CPU、XPU 上,以及 ROCm 上没开拒绝采样时,验证一律走 argmax,温度被忽略(eagle_utils.py#L720)。ROCm 会自动打开拒绝采样来避开这一条。
  • 多卡时各卡的 softmax 可能有细微差异,采样结果从 0 号卡广播。

数学验证:拒绝采样不改分布

这是数学验证,不是引擎实测:用纯 Python 标准库在一个玩具分布上按上面几种规则各抽 20 万次,看输出 token 的频率是否等于目标分布。目标 p 和草稿 q 是 6 个 token 上手写的两个分布:

P = [0.40, 0.25, 0.15, 0.10, 0.07, 0.03]   # 目标
Q = [0.20, 0.35, 0.05, 0.25, 0.10, 0.05]   # 草稿

规则对应关系:

规则对应源码
标准拒绝采样:x ~ q,u * q(x) < p(x) 接受,否则从 max(p−q, 0) 采vLLM _rejection_kernel + _resample_kernel;SGLang 打开拒绝采样后
one-hot 草稿:x = argmax q,u < p(x) 接受,否则从清零 x 的 p 采vLLM 默认 draft_sample_method="greedy"
只看目标、3 个兄弟:q 的前 3 个候选,一枚硬币累加 pSGLang 默认 TreeSpeculativeSamplingTargetOnly
同上但 threshold_acc = 0.5SGLang 调小阈值
对照组:直接用草稿采样的结果不做验证

另外在一个 4 个 token 的玩具马尔可夫链上,每轮提 2 个草稿,用标准拒绝采样生成长度为 3 的序列,和精确的联合分布(64 种结果)比。

脚本 verify_dist.py(约 220 行),运行 python3 -I verify_dist.py 200000 20261031,单核约 1.2 秒。种子 20261031 的输出:

规则χ²自由度p 值总变差距离接受率理论接受率
标准拒绝采样6.2950.280.00220.7018Σ min(p, q) = 0.70
one-hot 草稿4.3150.510.00170.2495p(argmax q) = 0.25
只看目标、3 个兄弟6.2050.290.00260.7510前 3 个候选的 p 之和 = 0.75
只看目标、threshold_acc = 0.5123688500.3471.0—
对照:直接用草稿91310500.3001.0—
3 token 序列、每轮 2 个草稿47.22630.930.0052——

换种子 7 再跑一遍,前三行和最后一行的 p 值是 0.36、0.48、0.75、0.90,两行反例的 χ² 仍在 9 万以上。

怎么读:

  • 三种无损规则的频率和目标分布的差距在抽样误差之内(p 值都不小,总变差距离在千分之几);两行反例的 χ² 大了几万倍,说明这个检验分得出「改了分布」和「没改分布」。
  • 同一对 (p, q),one-hot 的接受率只有 0.25,标准拒绝采样是 0.70:这就是默认不用草稿概率的代价。具体模型上差多少,待实测(C3-5)。
  • 只看目标的树规则接受率最高(0.75),因为它一次试 3 个候选;代价是验证的 token 更多。
  • 阈值调到 0.5 后接受率到 1.0,但输出分布偏了 0.35 的总变差距离。

与调度、KV 的交互

vLLM:草稿是「新 token」,被拒就回退计数

调度器把请求的 spec_token_ids 当作本步新 token 的一部分排进去,记在 scheduled_spec_decode_tokens 里(scheduler.py#L872),它们和其他 token 一起占 token 预算。分配 KV 时多带一个 num_lookahead_tokens(scheduler.py#L781):EAGLE 系和草稿模型要在目标查询范围之外再写 k 个位置的草稿 KV,所以预留 k 个槽(vllm.py#L615)。

vllm/v1/core/kv_cache_manager.py

        num_tokens_need_slot = min(
            num_tokens_main_model + num_lookahead_tokens, self.max_model_len
        )
        # ...
        num_tokens_to_cache = min(
            total_computed_tokens + num_new_tokens,
            request.num_tokens,
        )
        self.coordinator.cache_blocks(request, num_tokens_to_cache)

第二段是前缀缓存的保护:只有已确认的 token 所在的满块才算哈希、进缓存,没验证的草稿不会被别的请求命中(块哈希见 前缀缓存)。

验证结果回来后,被拒的个数从已算 token 数里扣掉:

vllm/v1/core/sched/scheduler.py

                if not output_is_stale:
                    if request.num_computed_tokens > 0:
                        request.num_computed_tokens -= num_rejected
                    if request.num_output_placeholders > 0:
                        request.num_output_placeholders -= num_rejected

块不还回去。被拒位置的槽留在请求的块表里,下一步从回退后的位置接着写,直接覆盖。EAGLE 系和草稿模型的注意力层有自己的 KV(uses_draft_kv_cache,speculative.py#L1966),和目标的 KV 分同一份显存,开投机后能放的块会少一些,少多少待实测(C3-5)。

SGLang:按最坏情况预留,结束时归还

SGLang 每个请求有两个长度:kv_committed_len(已确认)和 kv_allocated_len(已分配,schedule_batch.py#L937)。每个解码步前,eagle_prepare_for_decode 让已分配长度至少到「已确认 + 预留」:

python/sglang/srt/mem_cache/allocation_sizing.py

def get_alloc_len_per_decode() -> int:
    # ...
    # Spec decoding allocates max(topk * num_steps, num_draft_tokens) per decode step.
    # ...
    if page_size == 1 or spec_topk == 1 or not spec_algo.has_draft_kv():
        return max(spec_steps * spec_topk, spec_tokens)

预留量再乘 2(allocation_sizing.py#L70),注释说这是给重叠调度的双缓冲:调度器排下一步时,上一步的确认长度还没回来。按 Llama 的默认 (5, 4, 8),每个请求每步要预留 2 × max(20, 8) = 40 个槽位;按 (3, 1, 4) 是 2 × 4 = 8 个。这两个数是按代码算出的(估算),实际分配还会按页对齐。

验证完,已确认长度加上本轮接受的 token 数(草稿 + 补的一个,batch_result_processor.py#L825)。topk > 1 时被接受的路径在树里不一定连续,要把它们的 KV 挪到前面连续的位置(spec_utils.py#L788)。请求结束时,[已确认, 已分配) 这一段由 _release_overallocated_kv_indices 还给分配器(common.py#L414)。草稿模型的 KV 池和目标共用同一个槽位分配器和 req_to_token 表(eagle_worker_v2.py#L223)。

flowchart TB
  subgraph vllm_side ["vLLM"]
    v1["调度:1 + k 个新 token<br/>allocate_slots 多留 k 个槽"] --> v2["只给已确认 token 的满块算哈希"]
    v2 --> v3["验证后 num_computed_tokens 减去被拒数"]
    v3 --> v4["块留着,下一步覆盖"]
  end
  subgraph sglang_side ["SGLang"]
    s1["每步预留 2 × max(步数 × topk, 草稿数)"] --> s2["验证后 committed 加上接受数"]
    s2 --> s3["topk > 1 时把接受路径的 KV 挪到连续位置"]
    s3 --> s4["请求结束时归还 committed 到 allocated 之间的槽"]
  end

两边的共同点:被拒草稿的 KV 不进前缀缓存,也不单独释放,只是「确认长度」不往前走。不同点:vLLM 每步按需多留 k 个槽;SGLang 按最坏情况多留,换来重叠调度时不用等上一步的结果。

对照:mistral.rs

纯 Rust 的 mistral.rs v0.9.5 也做投机解码,配置只有 Off 和 Mtp 两种(config.rs#L19),草稿采样同样有 Greedy / Probabilistic 之分。验证用的是同一个标准规则:

mistralrs-core/src/speculative/verifier.rs

        let accept_prob = if q_i <= 0.0 {
            if p_i > 0.0 {
                1.0
            } else {
                0.0
            }
        } else {
            (p_i / q_i).min(1.0)
        };
        // ...
        let mut adjusted_probs = target_probs.sampling.clone();
        candidate_probs.subtract_from(&mut adjusted_probs)?;
        for prob in &mut adjusted_probs {
            *prob = prob.max(0.0);
        }

区别在位置:vLLM 和 SGLang 把这一步写在 Triton / CUDA kernel 里,一个 GPU block 处理一个请求;mistral.rs 上面这段是逐 token 的 Rust 循环,概率放在 Vec 里;开 cuda 特性时另有一条 GPU 上的稀疏拒绝采样路径。前向用 candle。

待实测

  • 接受率与平均每轮接受长度:n-gram、EAGLE3、草稿模型,贪心和温度 0.7 各一组:待实测(C3-5)。
  • 加速比:同一 8B 模型开 / 关投机,并发 1、8、48 的 TPOT 和总吞吐;SGLang 开投机后 max_running_requests 被设成 48,高并发时要单独对比:待实测(C3-5)。
  • vLLM draft_sample_method 取 greedy 和 probabilistic 的接受率差别,以及后者多占的显存:待实测(C3-5)。
  • 贪心下开关投机前后输出逐 token 是否一致:待实测(C3-5)。

小结

  • 投机解码用一次目标前向验证 k 个草稿,一轮出 1 到 k+1 个 token;草稿来自 n-gram 查表、独立小模型、EAGLE 头或 MTP 层。
  • vLLM 用 --speculative-config JSON,默认 V2 model runner 的 speculator;SGLang 用 --speculative-algorithm,EAGLE 系的步数、topk、草稿数按架构自动挑。两边默认都不开。
  • 接受规则是拒绝采样:min(1, p/q) 接受,被拒从 max(0, p−q) 重采,输出分布等于目标分布。两边默认都不用草稿概率(vLLM one-hot、SGLang 只看目标),仍然无损,接受率是 p(x) 或候选的 p 之和。
  • SGLang 的两个接受阈值默认 1.0,调小就不再无损;SGLang 的惩罚项、CPU / XPU 上的验证也和严格等价有出入。
  • 草稿占 KV 槽:vLLM 多留 k 个槽,被拒就回退已算 token 数;SGLang 每步预留 2 × max(步数 × topk, 草稿数),请求结束时归还。SGLang 开投机后并发上限默认 48。
  • 玩具分布上的数学验证(20 万次抽样)确认三种无损规则的频率与目标一致、两种放宽规则明显偏离。接受率和加速比要 GPU,待实测(C3-5)。

💬 关联讨论

发起讨论
加载中...