DeepInfer 源码完全拆解（八）：采样、惩罚与确定性校验——从 logits 到 token 的三重闭环

输入参数包：SamplerInputs 的分层

采样的参数不是散着传的，而是打包成 SamplerInputs，其内部再分成 driver / decode / prefill / speculative 几组：

这里有一个容易踩的概念坑：CUDA Graph 的输出张量带着 padding。开启 CUDA Graph 之后，hidden_states 的行数被对齐到图捕获时的固定形状，实际有效的行只占其中一部分。因此源码里必须先做一次 gather 把有效行摘出来，再算 logits：

1
hidden_states = hidden_states.move()
2
selected_token_indices = sampler_inputs.selected_token_indices
3
...
4
# CUDA Graph model_output retains graph-batch padding; gather trims hidden_states.
5
hidden_states = gather(hidden_states, selected_token_indices)
6
logits = get_logits(hidden_states, embedding, self.vocab_size)
7
del hidden_states

del hidden_states 这句不是随手写的。logits 的形状是 [total_q, vocab_size]，以 $V=129280$、total_q=1536 计算，fp32 的 logits 就是 $1536\times 129280\times 4 \approx 794\ \text{MiB}$。在显存本已紧张的推理服务里，能早一刻释放就早一刻释放，这也是后面所有 del logits 重复出现的原因。

driver / follower 分工

forward() 里有一处分支值得单独讲：

1
if driver_inputs is None:
2
    ...
3
    if self.max_MTP_value > 0 and not is_prompt:
4
        mtp_module.decode_forward(...)
5
    return SamplerOutput(post_funcs=[])

driver_inputs is None 意味着当前 rank 不是采样 driver。非 driver 的 rank 在这里只做两件事：参与 gather 的集合通信（gather 在前向里是一次跨 rank 收集），以及驱动本地的草稿头前向。它不物化 logprobs，也不采样。这是上一节”采样收拢到单 rank”的代码级体现。

prefill 路径：采样 + prompt logprobs

prefill 阶段的采样逻辑集中在 _prefill_golden_gen()：

1
logprobs, token_index, sample_results, _, _, _ = process_sampling(
2
    logits, driver_inputs.kernel, self.force_topk,
3
)

注意这里 process_sampling 只传了前三个参数，后面四个（is_verify、last_nextn_token_logprobs 等）都是 None，因此不会触发拒绝采样分支——prefill 阶段没有草稿需要验证。

prefill 路径有三件事发生在采样之前：

• token 黑名单。DEEPINFER_DECODE_TOKEN_BLACKLIST 环境变量指定一批永远不该被采出来的 token id，实现在 forward() 开头：

  1
  if self.decode_token_blacklist is not None:
  2
      logits.index_fill_(-1, self.decode_token_blacklist, -torch.inf)

  这是一个运维向的开关，用于在特定部署里硬性封禁某些 token（例如已知会引发下游系统故障的特殊控制符）。用 index_fill_ 而不是逐行 mask，是因为黑名单对整个 batch 是一致的。

• 首 token 抑制 EOS。first_n_token_no_eos 让生成的第一个 token 不能是 EOS：

  1
  for row, token_id in zip(prefill_inputs.eos_suppression_rows, prefill_inputs.eos_token_ids, strict=True):
  2
      logits[row, token_id].fill_(-torch.inf)

  这里只对”需要抑制”的那些行做原地写，所以先由 preparation 阶段把行号算好，采样时不判断。

• 结构化输出掩码：XGrammarUtils.apply_vocab_mask_by_step_inplace(logits, so_bitmasks[0], so_indices[0])。

顺序是有讲究的：黑名单是全局的，EOS 抑制是 per-row 的，语法掩码也是 per-row 的，三者都必须在 softmax 之前落下去，否则被 mask 掉的 token 会被算进归一化分母，导致采样出的 token 概率被稀释。

decode 路径：_verification()

decode 阶段无论是否开启 MTP，都走 _verification()。这个函数名字容易误解——它不只是”验证草稿”，而是整个 decode 采样流程：把 logits 按请求与步数重新映射、处理结构化掩码、跑 process_sampling、再收集 token logprobs / top-logprobs / sum_smaller_probs，最后汇总为 SampledInfo。

1
ftp = decode_inputs.mtp_first_token_positions
2
seqlen_q = decode_inputs.mtp_seqlens_q
3
row_to_req, local_step, is_verify = prepare_verification_metadata(ftp, seqlen_q, total_q, self.max_MTP_value)

prepare_verification_metadata 做的是三件事：把 total_q 行映射回请求（row_to_req）、给出每行在请求内部的第几步（local_step）、标出哪些行是”需要检验草稿”的行（is_verify）。在带 MTP 的 decode 里，一个请求的 seqlen_q 可能是 $1 + n$（1 个真实 token 加 $n$ 个草稿 token），布局是变长的——这正是第五篇讲的 DSpark 变长验证在采样侧的投影。

图源：Fast Inference from Transformers via Speculative Decoding 论文 Figure 1。每一行是一次迭代：草稿模型连续给出若干 token，目标模型一次前向并行算出这些位置的真实分布，然后按接受规则决定接受多少个、以及在第一个被拒绝的位置从哪里重新采样。DeepInfer 把”验证”这一列压进了采样内核：is_verify 标记出哪些行属于草稿位置，接受判定与残差重采样都在同一趟 logits 遍历里完成。

调用入口

1
def fused_sample(
2
    logits, sampling_inputs, *, tag="default", is_verify=None,
3
    logprobs_out=None, last_nextn_token_logprobs=None,
4
    last_next_n_token_ids=None, mtp_info_read_rows=None,
5
):
6
    assert logits.stride(1) == 1
7
    ...
8
    kernel = ops.fused_sample_tl(vocab_size, do_sample, do_reject_sampling)
9
    kernel(
10
        logits, transform_inputs.temperatures, transform_inputs.top_ps, transform_inputs.top_ks,
11
        sample_positions, seq_seeds, is_verify, logprobs,
12
        candidate_token_ids, last_nextn_token_logprobs, mtp_info_read_rows,
13
        last_next_n_token_ids, accept_result_verify, draft_selected_logprobs,
14
        accept_probs, seed_offset, lse_workspace,
15
    )

几处细节：

• assert logits.stride(1) == 1 强制最后一维连续。融合内核按行顺序扫描词表，只有保证 stride 为 1 才能做向量化读取与分段归约。
• 内核按 (vocab_size, do_sample, do_reject_sampling) 三元组特化编译。vocab_size 参与编译期特化是因为分段归约的粒度需要静态确定；后两个布尔量决定是否生成采样与验证的那些分支——编译期剪枝掉不用的分支，而不是运行时判断，这是保持内核紧凑的关键。
• lse_workspace 是一个 [total_q, 2] 的小 workspace。分块归约需要一个中间存放 log-sum-exp 的位置，用 fp32 保存部分和与部分最大值。
• seed_offset 用来分隔 RNG 命名空间，下一节细讲。

内部七个阶段

内核内部的顺序，从参考实现 fused_sample_ref 可以逐行对上：

1
logits.mul_((1.0 / transform_inputs.temperatures).unsqueeze(1))
2
apply_top_p_top_k_ref(logits, transform_inputs.top_ps, transform_inputs.top_ks)
3
logprobs = torch.log_softmax(logits, dim=-1)

第一步是温度缩放，用的是乘法而不是除法：mul_(1.0 / temperature)。这里有一个数值上的考虑——一次倒数加一次乘法，比每行做一次除法要快，而且避免了对每个元素做除法。代价是温度值本身被求了倒数，引入一次舍入。温度本身会被 clamp 到 SAMPLING_EPS = 1e-5 以上（见 SamplingParams），所以不会出现除以零。

一个容易忽略的细节：温度小于 SAMPLING_EPS 时，SamplingParams 不是只把温度夹住，而是同时把 top_p 设为 1.0、top_k 设为 -1：

1
if self.temperature < SAMPLING_EPS:
2
    self.temperature = SAMPLING_EPS
3
    self.top_p = 1.0
4
    self.top_k = -1

这是”贪心采样”的退化路径。温度趋零时 softmax 会退化为 argmax，此时任何 top-p/top-k 截断都只会引入不必要的排序开销——反正结果都是最大 logit 那个 token，不如直接放开截断，让内核走最便宜的路径。

第二步是 top-p 筛选。参考实现是这样一个稳定的前缀选择：

1
probs = torch.softmax(logits, dim=-1)
2
sorted_probs, sorted_indices = torch.sort(probs, dim=-1, descending=True, stable=True)
3
sorted_remove = sorted_probs.cumsum(dim=-1) - sorted_probs > top_p.unsqueeze(1)

用数学写出来：把概率降序排列为 $p_{(1)} \ge p_{(2)} \ge \cdots$，保留满足

的 token $i$。注意 cumsum - sorted_probs 这个操作——它是”前 $i-1$ 项之和”，等价于把第 $i$ 项排除在累积和之外再与阈值比较。这个写法比”累积和超过 p 就截断”宽松一格，保留了第一个使累积和越过阈值的 token，与 HuggingFace 的 TopPLogitsWarper 语义一致。改用别的口径（例如严格 cumsum > p）会让同一个 top_p 值在两套实现里选出不同的候选集，进而让同一模型在迁移时输出不一致。

stable=True 同样不是装饰：词表里大量 token 的 logit 完全相等（尤其是 logits 被大量 mask 成 -inf 之后），稳定排序保证并列元素按索引升序排列，使候选集的边界是确定的。后面 3.9 节会看到，DeepInfer 专门为这个性质写了测试用例。

第三步是 top-k 筛选，同样是稳定的：

1
order = torch.argsort(logits, dim=-1, descending=True, stable=True)
2
ranks = torch.empty_like(order)
3
rank_values = torch.arange(vocab_size, device=logits.device).expand_as(order)
4
ranks.scatter_(1, order, rank_values)
5
remove = (top_k >= 0).unsqueeze(1) & (ranks >= top_k.clamp(max=vocab_size).unsqueeze(1))

注意 top_k = -1 表示关闭，top_k >= 0 的条件判断就是为此。而 clamp(max=vocab_size) 处理的是 top_k 大于词表的情况——测试里专门构造过 max_top_k = vocab_size + 1024 的参数，实现必须能容忍。

top-p 与 top-k 的先后顺序固定为先 p 后 k。这与 vLLM 一致。顺序有影响：先做 top-p 会把候选集缩到 $k'$ 个，再按这个较小的集合做 top-k 排名，与先 top-k 再 top-p 的结果在边界上可能不同。

第四步是 log-softmax：

其中 $z_i$ 是经过温度缩放与掩码后的 logit。实现用分块的方式算 $\log\sum_j e^{z_j}$：因为词表有 129280 项，一行放不进 shared memory，必须切成若干段，每段内部求最大值与指数和，再跨段合并。lse_workspace 的 [total_q, 2] 就是为此预留的。合并的公式是标准的在线 log-sum-exp：

这个递归保证数值稳定——即使某段的 logit 全部是 1e30，也不会溢出，因为最大值先被提出来了。

第五步是随机采样，第六步是投机验证，第七步是拒绝重采样。这三步合起来构成本篇后面几节的主题。

为什么 do_sample=False 也要跑内核

fused_sample 有一个看起来奇怪的能力：sampling_inputs 可以只传 LogitsTransformInputs（温度、top-p、top-k），不传 seq_seeds 与 sample_positions，这时 do_sample=False，内核只算 logprobs 不采样。

这个路径是为 prompt logprobs 服务的。客户端如果请求了 prompt_logprobs，服务端需要对 prompt 的每一个位置给出该位置实际 token 的 logprob，这时需要一个”只算分布、不选 token”的算子。复用同一个内核、只是关掉采样分支，可以让”prompt logprob 的分布”与”采样的分布”在数值上完全一致——如果分成两套实现，用户就会看到一个诡异的现象：某 token 的 logprob 显示是 -1.2，但采样却从来没采到过它。

PyTorch RNG 的两个问题

如果直接用 torch.multinomial(probs, 1)，会立刻遇到两个问题。

问题一：结果依赖 batch 组成。torch.multinomial 内部按行取随机数，随机流的状态取决于这一行在整个 batch 中的位置。一个请求单独进来时排在行 0，与其他 200 个请求一起进来时可能排在行 137——两者会采出不同的 token。对于一个缓存了 KV 前缀、期望”相同输入得到相同输出”的服务来说，这是不可接受的：用户会发现同一段 prompt 在低峰期和高峰期给出不同答案。

问题二：结果依赖调用次数。全局 RNG 状态随采样次数推进，任何一次额外的采样（比如别的请求、或者服务端内部的探测性推理）都会改变后续结果。

DeepInfer 的解法是把随机性完全从”流状态”里抽出来，改为纯函数式的确定性映射：给定 (序列种子, 采样位置, 用途)，唯一确定一个均匀随机数。

种子与位置的构造

SamplingParams 里，每条序列在创建时分配一个 64 位种子：

1
@staticmethod
2
def _make_seed(seed: int | None) -> int:
3
    if seed is None:
4
        return random.getrandbits(64)
5
    if isinstance(seed, int):
6
        return seed & 0xFFFF_FFFF_FFFF_FFFF

seq_seed 允许客户端显式指定（用于复现），缺省时随机生成。种子是 per-sequence 的，不是 per-request：同一序列的两步 decode 用同一个种子，靠 sample_positions 区分。

采样位置来自哪里？看 MTPSamplingInput.build_padding：

1
sample_positions=torch.zeros(num_reqs, dtype=torch.int32, device="cuda"),

以及在测试里构造的：

1
sample_positions = torch.randint(0, 10**8, (b,), dtype=torch.int32, device="cuda")

位置在实现里是一个与序列位置绑定的整数。于是随机数由 $(seed, position)$ 这对值决定。

SplitMix64：从整数到均匀分布

把整数对映射成 $[0,1)$ 上的均匀数，DeepInfer 用的是 SplitMix64 混合函数。参考实现 _hash_to_uniform：

1
def _hash_to_uniform(seed: torch.Tensor) -> torch.Tensor:
2
    """SplitMix64 hash of seed -> uniform f64 in [0, 1); int64 ops emulate uint64."""
3
    z = seed.to(torch.int64)
4
    z = z + i64(0x9E3779B97F4A7C15)
5
    z = (z ^ urshift(z, 30)) * i64(0xBF58476D1CE4E5B9)
6
    z = (z ^ urshift(z, 27)) * i64(0x94D049BB133111EB)
7
    z = z ^ urshift(z, 31)
8
    return urshift(z, 11).to(torch.float64) * (1.0 / (1 << 53))

这段代码有三个值得展开的点。

为什么用整数哈希而不是 rand()？ 因为哈希是无状态的纯函数：输入相同的 $(seed, position)$ 必然得到相同的输出，与调用的先后、batch 的大小、线程数全都无关。这正是批次不变性需要的性质。

SplitMix64 的三个魔数。0x9E3779B97F4A7C15 是黄金比例 $\varphi$ 的小数部分乘以 $2^{64}$：

它的作用是让相邻整数在乘以它之后散开到整个 64 位空间的不同区域——这是 Weyl 序列的经典应用。后面两个常数 0xBF58476D1CE4E5B9 与 0x94D049BB133111EB 是 SplitMix64 的固定混合因子，配合右移 30、27、31 位形成 avalanche 结构：输入的任意一位变化都能扩散到输出的几乎所有位。

为什么要手写 urshift？ 因为 PyTorch 没有无符号 64 位算术，而 SplitMix64 的定义域是 uint64。实现用有符号 int64 模拟：

1
def urshift(x: torch.Tensor, n: int) -> torch.Tensor:
2
    shifted = torch.div(x, 1 << n, rounding_mode="floor")
3
    return shifted + (x < 0).to(torch.int64) * (1 << (64 - n))

逻辑右移 $x \gg n$ 在无符号语义下就是 $\lfloor x / 2^n \rfloor$；对有符号表示，负数需要一个补偿项 $2^{64-n}$ 把符号扩展位填回来。i64() 辅助函数则把大于 $2^{63}$ 的常量重新解释为负数。

最后一步为何是 urshift(z, 11) * 2^{-53}：取 64 位结果的高 53 位。这是因为 IEEE-754 双精度浮点的尾数恰好 53 位，取高 53 位再乘以 $2^{-53}$，得到的是 $[0,1)$ 上均匀可表示的数，且不引入任何舍入误差。这是一个把整数均匀性无损映射到浮点均匀性的标准技巧。

位置如何参与：三种互不相关的采样流

随机数不能只在种子和位置上哈希，还必须区分用途。同一个位置，采样目标模型的 token 和采样草稿 token 必须是两个独立的随机数，否则两者的误差会相关，破坏投机解码的无偏性。参考实现里：

1
sample_case = 1 if is_nextn_predict else 0
2
seq = (sample_positions.to(torch.int64) * 3 + sample_case) * i64(0x9E3779B97F4A7C15)
3
combined = seq_seeds.to(torch.int64) + (seq ^ (seq * (1 << 13)))
4
u = _hash_to_uniform(combined).unsqueeze(-1)

位置被乘以 3 再加 case，case 有三个取值空间：

在 fused_sample_ref 里可以清楚看到第三种的用法：

1
verify_seq = sample_positions.to(torch.int64) * 3 + 2
2
verify_combined = seed_block ^ (verify_seq * i64(0x9E3779B97F4A7C15))
3
U = _hash_to_uniform(verify_combined).to(accept_lp.dtype)

为什么接受判定要单独用一个随机数，不能复用采样那个？ 因为接受判定 $U \sim \text{Uniform}(0,1)$ 必须与目标分布的采样结果独立。如果两者用了同一个随机数，接收规则就不再产生正确的目标分布——这是投机解码无偏性证明的前提。乘 3 分流是最省的做法：三个用途各占一个位置槽，互不碰撞。

另外还有一层 seed offset 把”草稿”和”目标”分开：

1
_MTP_SEED_OFFSET = 1000000007
2
seed_offset = 0 if tag == "default" else _MTP_SEED_OFFSET

tag="mtp" 时种子整体偏移 $10^9+7$（一个质数）。这比”用同一个种子但不同 case”更强：它保证草稿采样与目标采样连哈希输入都落在完全不同的区域。

批次不变性：从代码性质到测试

把上面三条合在一起，random_sampling 的确定性就是一条代码级定理：

输出只依赖 $(\text{probs row}, \text{seq\_seed}, \text{sample\_position}, \text{is\_nextn\_predict})$ 四个量，与行在 batch 中的下标、batch 的总行数、其他行的内容都无关。

测试直接验证了这一点：

1
ref = random_sampling(probs, seq_seeds.to(torch.uint64), positions, is_nextn_predict=False)
2

3
perm = torch.randperm(n, device="cuda")
4
out_perm = random_sampling(
5
    probs[perm].contiguous(), seq_seeds[perm].to(torch.uint64), positions[perm].contiguous(), is_nextn_predict=False
6
)
7
assert torch.equal(out_perm, ref[perm])
8

9
for i in range(n):
10
    single = random_sampling(probs[i:i+1].contiguous(), seq_seeds[i:i+1].to(torch.uint64),
11
                             positions[i:i+1].contiguous(), is_nextn_predict=False)
12
    assert single.item() == ref[i].item(), f"row {i}: single={single.item()} != batched={ref[i].item()}"

三种情况：原序、随机置换、逐行单独跑，结果必须逐元素相同。第二段循环跑 $n$ 次单行 kernel，是最严格的一种——它把”batch 化不改变结果”这个命题变成了可证伪的断言。

这条性质在服务层的价值。有了批次不变性，才会出现这些行为：

• 同一请求在不同并发下产生完全相同的输出，前缀缓存与投机解码的可复现性才有基础。
• 测试可以用小 batch 跑出的结果作为 golden 基准，去验证大 batch 的正确性——否则正确性验证只能靠统计。
• 出问题时可以用单行重放来定位，不必复现整个 batch 的组成。

这也解释了为什么第六篇里 batch-invariant 会作为一个独立的工程目标出现——它不是采样的附属性质，而是整个引擎的一条系统级不变量。

图源：Thinking Machines Lab, Defeating Nondeterminism in LLM Inference。这组图讲的是更底层的一层非确定性：同一个请求与别的请求一起批处理时，因为归约顺序、分块大小随 batch 变化，数值结果会与单独处理时不同。DeepInfer 在采样这一层用哈希随机数切断的是另一条耦合——随机流的状态——但两件事的方向一致：让”结果”只取决于请求自身，而不是取决于它和谁一起跑。

分布算术

设目标分布为 $p$，草稿分布为 $q$，当前候选草稿 token 为 $x$。DeepInfer 实现的是拒绝采样（rejection sampling）而非 vLLM 常用的贪心比较，两者的接受判定不同：

• 贪心比较：$x$ 与 $\arg\max p$ 相同就接受——只在温度趋零时等价于 $p$。
• 拒绝采样：以 $\min(1, p(x)/q(x))$ 的概率接受——对任意 $p$ 都成立。

DeepInfer 采用后者，因为服务端允许任意温度。实现里对应的是这一段：

1
draft_token_lp = draft_logprobs.gather(1, safe_tids.unsqueeze(1)).squeeze(1)
2
draft_selected_logprobs = draft_token_lp.masked_fill(sentinel_mask, float("-inf"))
3
target_lp = logprobs.gather(1, safe_tids.unsqueeze(1)).squeeze(1)
4
accept_lp = torch.clamp(target_lp - draft_token_lp, max=0)
5
U = _hash_to_uniform(verify_combined).to(accept_lp.dtype)
6
accepted = (accept_lp.exp() > U).masked_fill_(sentinel_mask, False)

逐步翻译：

1. target_lp - draft_token_lp 是 $\log p(x) - \log q(x)$。在 fp32 里做减法而不是除法，是为了避免 $q(x)$ 极小（例如 e-30）时的下溢。
2. clamp(max=0) 把比值封顶到 1：$\exp(\min(\log p - \log q, 0)) = \min(1, p/q)$。这一个 clamp 就是整个接受概率公式。
3. accept_lp.exp() > U 实现”以概率 $a$ 接受”——标准做法，$U$ 是独立的均匀随机数。
4. sentinel_mask 把草稿 token 为 -1 的行强制置为拒绝。哨兵值用于填充位置与”这一轮没有草稿”的请求。

残差分布。被拒绝后需要从

采样。实现里有一个关键偏离：

1
residual = (sample_probs - (draft_logprobs - DRAFT_PROB_PENALTY).exp()).relu_()

注意 DRAFT_PROB_PENALTY 这一项。草稿概率被整体下调之后再参与残差计算。为什么要加惩罚？

纯理论的残差 $\max(p - q, 0)$ 有一个实践问题：$q$ 是由草稿模型给出的估计值，可能与真实草稿分布有偏差；如果 $q$ 被系统性高估（例如草稿模型对自己的选择过度自信），残差会被低估，重采样的分布就偏离 $p$。对 $q$ 施加一个常数惩罚会让残差整体变大——也就是更保守地估计”多少概率质量需要被重新分配”，从而在 $q$ 不准时仍然接近目标分布。测试里的分布校验正是围绕这个惩罚项的量级展开的（见 6.2 节）。

relu_() 把负值归零，得到未归一化的残差权重。

接受概率的估计值 accept_probs：

1
accept_probs = (1.0 - residual.sum(dim=-1)).clamp_min_(0.0)
2
accept_probs = torch.where(is_verify, accept_probs, 1.0)

这个量不是用于判定的，而是回传给上层做投机深度控制的：$1 - \sum_i r_i$ 正是随机采样的接受概率 $\sum_i \min(p_i, q_i)$。第五篇讲过 DSpark 会依据这个估计值动态调整下一轮的投机深度——如果某些请求的接受率持续偏低，就该减少草稿 token 数，省下白算的验证开销。非验证行恒为 1.0。

拒绝后的重采样

1
reject_sample = random_sampling_ref(residual, seed_block, sample_positions, is_nextn_predict=True)
2
sample_result = torch.where(is_verify, torch.where(accept_verify, draft_tids, reject_sample), sample_result)

三层嵌套的 where 各自负责一件事：非验证行用原来的正常采样结果；验证行且接受则直接用草稿 token；验证行且拒绝则用残差重采样结果。is_nextn_predict=True 让重采样走 case=1 的随机流，与正常采样区分开。

内核侧的固定点采样

GPU 内核里的采样不用浮点 CDF 比较，而是定点数。测试注释里透露了这一点：

1
total_eps = h / float(1 << 31)

分母 $2^{31}$ 说明内核把概率量化成 31 位定点数，逐行累积后与哈希得到的均匀数比较。这带来一个上限：$V$ 个 token 的量化误差累计可达 $V/2^{31}$，$V=129280$ 时约 $6\times10^{-5}$。测试用 total_eps 作为容差断言 accept_probs_out 与参考值的偏差，正是因为这个量化的存在。

固定点采样的另一个后果是”零概率永不采中”。浮点 CDF 在极小概率下可能因舍入给出非零权重；定点数下，$p_i < 2^{-31}$ 的项会量化成 0，累积区间宽度为零，永远落在区间边界之外。测试里对这一点有专门断言：

1
sampled_p = probs.gather(1, out.unsqueeze(1)).squeeze(1)
2
assert (sampled_p > 0).all(), f"sampled a zero-probability token, min={sampled_p.min().item()}"

以及退化场景测试——只有单个非零概率时，必须采到那一个：

1
probs[0, 3] = 1.0
2
probs[1, 7] = 0.5
3
probs[2, 2] = 1e-9
4
probs[3, 5] = 1e-4
5
for seed in (1, 42, 1234):
6
    out = random_sampling(probs.clone(), _seeds(4, seed), _pos(4), is_nextn_predict=False)
7
    expected = torch.tensor([3, 7, 2, 5], dtype=torch.int32, device="cuda")
8
    assert torch.equal(out[:4], expected)

注意第 3 行 p = 1e-9。定点量化下 $1e\text{-}9 \ll 2^{-31} \approx 4.66\times10^{-10}$……实际上 $2^{-31}$ 是 $4.66\times10^{-10}$，$10^{-9}$ 略大于它，所以仍能被表示。测试特意挑了紧贴量化底线的取值。

与 vLLM、TensorRT-LLM 的对照

三个框架都实现了拒绝采样，但侧重不同：

DeepInfer 的两点差异都指向同一个目标：在草稿模型质量不可控的前提下保持无偏。惩罚项是对草稿估计误差的保险，接受率回传是让调度器自己去适应草稿质量。

站内关于 vLLM 投机解码的实现见 vLLM 源码完全拆解（十）：投机解码，TensorRT-LLM 侧见 TensorRT-LLM 源码完全拆解（十）：投机解码。

全局 token 黑名单

1
if blacklist := envs.DEEPINFER_DECODE_TOKEN_BLACKLIST:
2
    values = parse_values(blacklist, 0, self.vocab_size - 1, env_key="DEEPINFER_DECODE_TOKEN_BLACKLIST")
3
    self.decode_token_blacklist = torch.tensor(values, dtype=torch.long).cuda()

黑名单在 __init__ 时构造一次，之后每步用 index_fill_ 施加。选择 -inf 而不是把概率置零，是因为后面还有 log-softmax：置零的概率会在归一化时被其他 token 瓜分，看起来没问题，但当所有候选都被置零时会出现除以零；用 -inf 则统一表现为”这些 token 的 logprob 也是 -inf”，采样永远不会选中，且数值上与掩码语义一致。

golden token 覆盖

采样层支持一种强制输出指定 token 的模式：

1
sample_results_cur = (
2
    torch.where(driver_inputs.golden_output == -1, candidate_vocab_ids, driver_inputs.golden_output)
3
    if driver_inputs.golden_output is not None else candidate_vocab_ids
4
)

语义是”逐行覆盖，-1 表示不覆盖”。用途是 golden teacher forcing——让推理引擎在给定 token 序列上运行，从而与训练侧的数值对齐、或用于复现某个已知的输出。注意这发生在采样之后：内核照常算出采样结果，golden_output 只是在最后把它替换掉。这样做的好处是，即便某行被覆盖了，它的 logprobs 仍然是真实分布下的值（_update_golden_tokens 里会按覆盖后的 token 重新查 logprob），所以 golden 模式下的 trace 依然可以被用于数值调试。

结构化输出掩码的时序

结构化输出在采样侧的接口是”给定每行一个 bitmask，把非法 token 的 logit 置为 -inf”。第七篇讲过 bitmask 由 XGrammar matcher 产生，这里只看它在采样时序中的三个位置：

1. 掩码准备与 logits 计算重叠（上一节已看代码）。

2. 掩码按步施加。带 MTP 时，每个草稿步有自己的掩码：

   1
   if so_indices is not None:
   2
       for step, indices in enumerate(so_indices):
   3
           if len(indices):
   4
               apply_token_bitmask_by_step_inplace(logits, so_bitmasks[step], indices, ftp, step)

   step 从 0 到 max_MTP_value，对应”第一个真实 token”与各草稿 token 的位置。每一步的掩码是把该步之前的 token 喂进 matcher 后推进出来的状态，所以不能复用同一个掩码——这是第六篇讲的 matcher 回滚机制存在的原因。

3. 掩码在验证之前。注意 apply_token_bitmask_by_step_inplace 在 process_sampling 之前调用。也就是说，被语法禁止的 token 从候选集里就被删掉了，草稿 token 如果落在禁止集合里，其 logprob 会是 -inf，接受判定必然失败。这比”先验证再检查语法”更安全：后者会允许一个语法非法的草稿被接受，然后由下游发现并重试。

掩码分配的细节

1
if not self.bitmasks_cpu or self.bitmasks_cpu[0].shape[0] < num_sampling_token:
2
    self.bitmasks_cpu = [
3
        allocate_vocab_mask(num_sampling_token, self.vocab_size, pin_memory=True)
4
        for _ in range(self.max_MTP_value + 1)
5
    ]

pin_memory=True 是为了后续的异步 H2D 拷贝；数组长度 max_MTP_value + 1 是为每个草稿步各留一块。每块的大小是 $\text{num\_tokens}\times \lceil V/32 \rceil$ 字节——bitmask 按位压缩，$V=129280$ 时每行 4040 字节，比 510 KiB 的 float mask 小三个数量级。这个设计对采样阶段尤其重要：如果掩码是浮点张量，H2D 拷贝量会挤占采样带宽。

确定性：同输入必然同输出

1
logprobs_out2, candidate_out2, accept_out2, draft_selected_out2, accept_probs_out2 = fused_sample_fn(
2
    logits.clone(), sampling_params, tag=tag,
3
    is_verify=is_verify if do_reject_sampling else None, **draft_kwargs,
4
)
5
assert torch.equal(logprobs_out, logprobs_out2)
6
if do_sample:
7
    assert torch.equal(candidate_out, candidate_out2)

这里用的是 torch.equal 而不是 allclose——要求逐位相同，包括所有输出缓冲区。这条断言把一个强得多的承诺固化下来：融合内核的结果不是”近似确定”，而是”完全确定”。

拒绝采样的分布正确性

上面讲过残差重采样的分布必须回到 $p$。测试 _check_reject_sampling_distribution 是一段相当讲究的统计验证，以 batch=1048576×2、每行重复 1024 轮的量级采样，然后对两个 token 的计数做上界检验：

1
total_eps = 1.0 - math.exp(-DRAFT_PROB_PENALTY)
2
real_max_prob = target_probs * (1.0 + eps)
3
expected_counts = target_probs.double() * total_samples
4
max_expected_counts = real_max_prob.double() * total_samples
5
variance_prob = real_max_prob.clamp_max(1.0)
6
variance = (variance_prob * (1.0 - variance_prob)).double() * total_samples
7
...
8
log_factor = math.log(vocab_size / 1e-4)
9
tolerance = log_factor / 3.0 + (2.0 * variance * log_factor + log_factor**2 / 9.0).sqrt()
10
excess = actual_counts - max_expected_counts
11
assert (excess <= tolerance).all()

三个关键点。

上界不是 $p$ 而是 $p(1+\epsilon)$。因为残差计算里草稿概率被惩罚过，采样出的实际分布会有系统性偏移，$\epsilon = 1 - e^{-\text{PENALTY}}$ 正是这个偏移的上界。测试不是去验证”分布等于 $p$“——那是不可能的——而是验证”分布被 $p(1+\epsilon)$ 支配”，即落在可证明的范围之内。这是一个数学上诚实得多的断言。

容差是 Bernstein 界而不是正态近似。$t = \frac{\log \text{factor}}{3} + \sqrt{2\sigma^2\log\text{factor} + \frac{(\log\text{factor})^2}{9}}$ 是 Bernstein 不等式的形式。它比 Chernoff 界在方差小的时候更紧，且不需要假定正态——对于二项计数这种场景更合适。

边界用例覆盖到 1e-10 量级。测试参数表里有 16 组边界，包括 draft_equal_1e-10、draft_two_orders_below（草稿概率比目标低两个数量级）、rare_token0_draft_opposes（稀有 token 上草稿与目标完全对立）。这些用例是冲着浮点下溢去的：草稿概率 1e-10 时，$q(x)$ 在 fp32 里已经接近非规格化数区间，减法与指数的精度都会退化。

零残差的退化情形

1
logits[:, -2:] = 0.005  # Tied maxima in the final, partial chunk.
2
draft_logprobs = target_logprobs + math.log1p(-1e-6)
3
draft_logprobs[:, 0] = math.log(1e-6)
4
draft_ids = torch.arange(b, dtype=torch.int32, device="cuda") % 3 - 1
5
residual = (target_logprobs.exp() - (draft_logprobs - DRAFT_PROB_PENALTY).exp()).relu_()
6
assert (residual < 2**-31).all()

构造的是”残差小到定点量化下限”的情形：草稿分布比目标略小一点点（log1p(-1e-6)），惩罚之后残差落在 $2^{-31}$ 以下，定点化后变成 0。此时接受概率恰好是 1，所有草稿都该被接受——expected_accept = draft_ids > 0 验证了这个边界。

这段测试实际上守的是”量子化不会把应当接受的草稿误判为拒绝”。如果内核的定点化是简单截断而非舍入，残差可能被量化成一个极小的正值，accept_probs 变成 $1-10^{-9}$ 而不再精确等于 1，后续的统计就会漂移。

有符号零与并列排序

1
logits = torch.zeros((b, h), dtype=torch.float32, device="cuda")
2
mixed_zeros = logits.clone()
3
mixed_zeros[:, ::2] = -0.0
4
...
5
assert torch.equal(logprobs, expected_logprobs), "+0 and -0 must share the same stable tie order"
6
assert torch.equal(ids, expected_ids)

构造一个所有 logit 都相等的行，一半是 $+0$ 一半是 $-0$。IEEE-754 里 $+0 = -0$ 在数值比较上相等，但位模式不同；如果排序实现比较的是位模式（某些 SIMD 比较指令会区分），或者排序不稳定，就会让 $+0$ 和 $-0$ 排成两组，从而选出不同的候选前缀。

测试进一步验证保留的是一个连续前缀：

1
expected_prefix = torch.arange(h, device=logprobs.device).unsqueeze(0) < n_kept.unsqueeze(1)
2
assert torch.equal(finite, expected_prefix), "ties must retain a contiguous prefix in token order"

即 mask 保留的 token 必须是 $[0, n_{kept})$ 这样一个下标连续区间。这条性质对缓存与调试都很重要：如果并列元素散布在词表各处，同一个 top_p 在不同实现里会保留完全不同的集合。

top-p/top-k 的三条边界性质

_check_top_p_top_k 用三组断言刻画截断的正确性：

1
# 1) No pruned probability may exceed a surviving probability.
2
pruned_prob_lower = torch.nextafter(pruned_max, torch.full_like(pruned_max, float("-inf"))).exp()
3
survive_prob_upper = torch.nextafter(survive_min, torch.full_like(survive_min, float("inf"))).exp()
4
assert (pruned_prob_lower[has_pruned] <= survive_prob_upper[has_pruned]).all()
5

6
# 2) surviving probs equal renormalized full softmax over same set.
7
diff_sum = (prob - surv_norm).abs().sum(dim=-1)
8
assert (diff_sum <= threshold).all()
9

10
# 3) The selected prefix is minimal.
11
assert ((n_survive <= k_limit) & (prefix_mass <= top_p + eps)).all()
12
assert ((n_survive >= k_limit) | (selected_mass > top_p - eps)).all()

三条性质分别是：被剔除的概率不高于被保留的最小概率（截断发生在正确的位置）、保留集的相对分布与全量 softmax 一致（没有在排序过程中扰动概率）、保留集是最小的（去掉最小项就够不到 top_p，加上它就够到）。

第一条用了 nextafter 把边界放宽一个 ULP，注释解释了原因：Ascend 平台上比较的是 exp(FP32 logprob)，不同的 logits 可能因为指数函数的舍入而并列。这不是”近似搜索容差”，而是精确数值比较所必需的边界处理。

第三条的 eps 是 1e-4，与第二条的容差一致。

测试矩阵的规模

把参数组合列出来能看出这套测试的成本意识：

1
for b in [256, 1536]
2
for h in [1024, 129280]
3
for do_sample in [True, False]
4
for do_reject_sampling in [True, False]

$h$ 取两个值：1024 是一个放得进单次归约的小词表，129280 是 DeepSeek 系列的真实词表。129280 这个数不是 2 的幂，也不是 128 的整数倍（$129280 = 2^8 \times 505 = 1010\times 128$，实际是 $129280 = 1010 \times 128$，是 128 的倍数但分段数不整），所以分段归约的最后一个分段是残段——正是最容易出错的形状。测试注释里明确点出过这一点：

1
logits[:, -2:] = 0.005  # Tied maxima in the final, partial chunk.

分布测试的 batch 更是到 $1048576\times2$，Ascend 平台因为运行时长降到 65536。这个量级才能在 1e-10 概率上观察到足够的计数。

采样算子的融合程度

TensorRT-LLM 的”分档”（根据 top-k 大小选择不同 kernel）与 DeepInfer 的”编译期特化”（按 vocab_size/do_sample/do_reject_sampling 特化）是两种不同的应对策略：前者在运行时选择预先编译好的若干实现，后者在编译期把不需要的分支剪掉。前者更容易扩展新档位，后者在固定部署下更紧凑。TensorRT-LLM 侧的细节见 TensorRT-LLM 源码完全拆解（十七）：Sampler。

随机性的来源

这是 DeepInfer 与其他三家差别最大的一点。

Philox 与 SplitMix64 是同一类思路的两种实现：都用计数器/位置作为输入的无状态映射，而不是流式 RNG。区别在于 Philox 是分组密码式的 4×32 位结构，适合 SIMD 批量生成多个随机数；SplitMix64 是单流 64 位混合，更适合”每行只需要一个均匀数”的场景。DeepInfer 采样每行只消耗 1 到 3 个均匀数，SplitMix64 的开销更低。

投机验证的归属

DeepInfer 把验证放进采样内核，收益很直接：logits 只读一次。验证要用的目标分布 logprobs 就是采样算出来的那一份，不需要再算一遍。代价是内核复杂度上去了——fused_sample 的条件分支比纯采样版本多不少，这也是为什么需要 do_sample / do_reject_sampling 两个编译期开关把它切成四种形态。

llama.cpp 侧的投机验证在 llama.cpp 源码完全拆解（九）：投机解码 里有完整拆解。