DeepInfer 源码完全拆解（五）：投机解码——DSpark 变长验证、DFlash 草稿头与置信度调度

投机解码的原始形态

投机解码（speculative decoding）的出发点是一个很朴素的不对称：大模型的单次前向很贵，但一次前向可以对很多个位置同时算 logits。既然如此，与其每步只算 1 个 token 的分布，不如先找一个便宜的来源猜出接下来 $K$ 个 token（这叫草稿，draft），然后用大模型一次前向同时算出这 $K+1$ 个位置的分布，再逐个判断草稿猜得对不对。猜对了就一次前进多个 token，解码步数直接减少。

这项技术由 Leviathan 等人于 2022 年提出（arXiv:2211.17192，Fast Inference from Transformers via Speculative Decoding），Chen 等人同期给出了把验证写成「拒绝采样」形式的版本（arXiv:2302.04858，Accelerating Large Language Model Decoding with Speculative Sampling）。两者的共同点在于：草稿不改变输出分布。这不是近似，而是有严格证明的等价性——只要草稿的采样分布 $q$ 与目标分布 $p$ 满足特定的接受规则，最终输出的分布与直接用 $p$ 采样完全一致。这是投机解码能进生产的前提。

标准做法（以 Chen 等人的形式为准）是：草稿模型对第 $i$ 个位置给出分布 $q_i$，采样出草稿 token $x_i$；目标模型算出对应位置的 $p_i$。逐个位置用

以概率 $a_i$ 接受草稿 token。第一个被拒绝的位置，从残差分布

重新采样；这之后的所有草稿全部丢弃。全部接受时，再在最后一个位置多采一个全新 token（这一步叫 bonus token）。可以证明这样得到的样本服从 $p$，从而整个序列的分布与原始解码器一致。

为什么通用框架里这一步很别扭

把这套规则塞进一个现代推理引擎，会立刻撞上几个工程问题：

• 草稿模型本身要跑一次前向，这需要它自己的权重、自己的 KV 缓存、自己的注意力元数据。如果草稿模型是独立的小模型，那么引擎里就要维护两套 KV 缓存池和两套注意力后端，调度器要协调两次前向的顺序；
• 验证需要逐位置的 $p_i(x_i)$ 与 $q_i(x_i)$，也就是要对「草稿 token 在每个位置上的目标 logprob」做一次 gather。如果草稿不是自回归生成而是一次前向并行吐出多个位置，那么草稿侧还要保存自己每个位置的完整分布（或者 top-k 近似），供验证时反查；
• 接受结果要回传给调度器：本轮实际接受了几个 token、KV 缓存里那几行该保留、请求的序列长度加多少，都必须以设备上的张量形式表达出来，并且要和下一步的输入构造对齐。

EAGLE（arXiv:2401.15077）与 EAGLE-2、Medusa（arXiv:2401.10774）代表了另一条路线：草稿不由独立模型给出，而是由挂在目标模型上的额外头给出，并且草稿生成与目标模型共享隐状态，因此草稿的质量更高、开销更低。DeepInfer 的 DFlash 属于这一路线的变长 block 版本：一次前向并行吐出 dflash_block_size 个草稿位置。

DeepInfer 的三个选择

DeepInfer 在这套通用问题上做了三个明确的选择，构成本篇的三条主线：

第一，草稿与目标共享注意力计算的分工是「反向共享」。DFlash 不是让草稿头去看目标模型的 KV，而是让目标模型的特征反过来喂给草稿层：主模型 forward 时顺手把指定层的 hidden 投影成草稿 KV，写进草稿自己的缓存（update_main_cache），草稿层再基于这份缓存做块内因果注意力。这样草稿层不需要重跑一遍主干注意力。

第二，验证被压进一次 TileLang 融合 kernel。fused_sample 一个 kernel 同时完成温度缩放、top-p/top-k 过滤、目标采样、草稿 logprob 对齐、拒绝采样判定与残差重采样。整个验证过程的中间张量不出 kernel。

第三，额外加了一个「置信度验证器」。这不是经典投机解码的组件：DFlash 的 ConfidenceHead 对每个草稿位置预测一个接受概率，调度器拿这个预测决定下一步要不要继续投机、投机多深。它把「验证」从输出正确性的范畴扩展到了调度决策的范畴。

一句话概括 DeepInfer 的做法

用主模型的中间层特征喂一个轻量草稿头，一次前向并行吐出 $\gamma$ 个草稿与每个位置的置信度；验证时用拒绝采样逐位置判定，把「被接受的前缀 + 一个重采样 token」当作本轮输出，并用置信度决定下一步验证多深。

这套机制的完整形态来自 DeepSeek 团队 2026 年 7 月提交的 DSpark 论文，下一节先把论文口径交代清楚，再逐段拆代码。

STS 在数学上做什么

论文对置信度的定义值得单独拎出来，因为它是整个调度逻辑的基础。置信度头对每个草稿位置 $k$ 输出一个标量 $c_k \in (0,1)$，它不是「这个 token 被接受的无条件概率」，而是

即一个条件概率。于是由链式法则，一个长度为 $k$ 的草稿前缀整体存活（全部被接受）的概率是累乘：

这就是代码里 accept_rates.push_back(accept_rate) 那行累乘的来源（见后文）。为什么必须区分条件概率与联合概率？因为调度器要算的是「期望接受长度」$\tau$——它是「我多验证一个位置能多拿到多少 token」这个问题的答案，而 $\tau$ 的增量正是 $a_k$。用无条件概率直接算会系统性高估，因为后面位置的接受是以前面位置被接受为前提的。

STS 的做法是：把 $c_k$ 逐个做温度缩放 $c_k \mapsto \sigma(\log \mathrm{odds}(c_k)/T_k)$，从 $k=1$ 到 $k=\gamma$ 依次做一维网格搜索，每一步只调当前的 $T_k$，使累乘值 $a_k$ 的 ECE（Expected Calibration Error，期望校准误差）最小，前面已校准的位置固定不动。

温度缩放有两个性质让这件事成立：

• 保序：$T_k$ 是单调变换，不改变草稿 token 之间的相对排序。所以即便校准过的值变了，置信度头学到的「哪个草稿更好」的判断仍然保留——这一点对调度器很关键，因为它在做 top-K 选择时依赖排序。
• 只改幅度不改顺序：因此校准可以在草稿头训练完之后单独做，不需要重训。这是工程上的大便利。

也正因为 $T_k$ 按位置分别搜索，最终的 $T$ 是一个长度等于草稿数的列表。这解释了代码里那个看似奇怪的接口设计：

1
temperature = envs.DEEPINFER_DFLASH_CONFIDENCE_HEAD_TEMPERATURE
2
temperatures = temperature if isinstance(temperature, tuple) else (temperature,) * num_steps
3
self.temperature = tuple(max(value, SAMPLING_EPS) for value in temperatures)

环境变量的类型既可以是标量（未做 STS、统一温度）也可以是元组（做过 STS、每位一个温度），而 num_steps 就是草稿长度 max_MTP_value。max(value, SAMPLING_EPS) 是把温度夹到一个极小正数下界——温度趋近 0 会让 logits 缩放爆炸，SAMPLING_EPS 是采样代码里通用的防除零常数。

调度目标是全局吞吐，不是单位置接受率

论文把验证长度的选择写成了一个全局吞吐最大化问题（Algorithm 1），而不是简单的「置信度低于阈值就砍掉」。理由很直接：验证一个低置信度 token 的机会成本取决于系统当前负载——低并发时它几乎不占什么资源，高并发时它占的是别人也要用的 batch 容量。静态阈值（论文点名了 threshold-based verification heuristics）在这种场景下必然次优。

问题形式化如下。设有 $R$ 个活跃请求，请求 $r$ 的每位置置信度是 $c_{r,1},\dots,c_{r,\gamma}$，调度器要为每个请求选一个验证长度 $\ell_r \in \{0,\dots,\gamma\}$。目标是在给定目标模型 batch 容量 $K$（论文用 $K$ 表示这一轮能验证的 token 总数）下最大化

其中分子是期望产出的 token 数（各请求期望接受长度之和），分母是这一轮目标前向的时间。$\mathrm{Time}(\cdot)$ 不是解析式，而是实测出来的性能曲线 SPS(B)（samples per second as a function of batch size）——这就是论文所说的 profiled step curve 与 engine-specific throughput profiles。

这个形式化解释了为什么「置信度排序 + 全局 top-K」是对的而「逐请求阈值」是错的：在全局预算下，取所有 (请求, 位置) 对中 $a_k$ 最高的那批，等价于按边际收益排序花钱。这也正是 DeepInfer 代码里那个 top-K 选择的来源。

非预测性（non-anticipating）：一条正确性硬约束

论文里有一段容易被略过但极其重要的论证。无损投机解码要求非预测性：准入决策不能依赖于未来的候选 token。而 DSpark 的置信度头依赖「上一步已采样 token」的 Markov 特征，要算出 $a_{r,k+1}$ 就必须先知道 $x_{r,k}$ 这个实际被采出来的候选 token。如果调度器做一次「回顾式的全局搜索」（先看完整块再决定谁进），就会把 $x_{r,k}$ 泄漏进第 $k$ 步的准入决策，引入选择偏差（论文在附录 A 给了一个具体的反例）。

这条约束是硬性的——违反它就不再是无损投机解码了。它直接决定了工程实现必须流式下发决策，不能「先收齐再统一裁剪」。后面会看到，DeepInfer 用「延后两步」的方式实现了这一点：用两步之前的置信度只决定总容量 $K$，而候选 token 在当步仍然按最新的置信度排序——论文原话是 the selection mechanism is fundamentally rank-preserving: the most confident draft tokens are always prioritized for verification。近似的是容量，不是排序，这是把非预测性与 ZOS 兼容性调和起来的关键。

草稿头从哪里挂上去

DFlash 是 deepinfer/model_executor/models/mtp/dflash.py 里的一个 nn.Module，它挂在主干模型的 MTP 层位置上。构造时它做了一件很明确的事——从模型容器里把 MTP 层切出来：

1
num_non_mtp_layers = cfg.num_hidden_layers + 2  # 主模型层数 + emb 层 + lmhead 层
2
if len(model_modules) <= num_non_mtp_layers:
3
    raise IndexError(f"模型没有 MTP 层/层数不对, 总层数{len(model_modules)}, ...")
4
modules = model_modules[num_non_mtp_layers:]
5
if not all(isinstance(module, DFlashNextPredictHeadHC) for module in modules):
6
    raise NotImplementedError("Only DFlash speculative heads are supported")
7

8
self.MTP_num_layers = len(modules)
9
assert self.MTP_num_layers == nextn_predict

几个约束值得点出来：

• + 2 是因为模型容器里多了 embedding 层与 lm_head 层两个 entry。这是一个典型的「读代码时容易数错」的地方——层数不是 num_hidden_layers，而是它加 2。
• 只支持 DFlash 类型的头。DeepSeek-V3 的 MTP 头与 DFlash 头的结构不同，DeepInfer 明确只留 DFlash 这一条路径，遇到别的类型直接 NotImplementedError 而不是静默降级。这与第一篇讲的「只服务一个模型所以可以彻底特化」是同一个取向。
• MTP_num_layers == num_nextn_predict_layers 的断言把「模型配置的层数」与「实际拿到的模块数」锁死，任何一侧改动都会立即报错。

main cache：把主模型特征写进草稿侧 KV

DFlash 最核心的设计是 update_main_cache。它的 docstring 写得很准确：

1
Write every DFlash layer's main-token keys from the target-layer features;
2
runs inside the main model forward.

也就是说，主模型做前向时顺手把草稿 KV 写了。流程是：

1
main_hiddens = self.projection_norm(self.projection(dflash_input))
2
main_k = self.main_kv_proj(main_hiddens).unflatten(-1, (self.MTP_num_layers, self.qk_head_dim))
3
for i, layer in enumerate(self.layers):
4
    layer.update_main_cache(
5
        main_k=main_k[:, i],
6
        kv_cache=nextn_kv_caches[layer.layer_idx],
7
        main_positions=main_positions,
8
        slot_mapping_store_swa=input_metadata.slot_mapping_store_swa,
9
        mirror_cp_reverse_index=mirror_cp_reverse_index,
10
    )

逐项解释：

• self.target_layer_ids = cfg.dflash_cfg.dspark_target_layer_ids 是要采集哪几层主模型特征的配置。projection 的输入维度是 cfg.hidden_size * num_target_feature_layers——也就是说它把若干层的 hidden 拼接后投影回一个 hidden size。这是 EAGLE 系「用目标模型中间层特征喂草稿头」的同一路数，DFlash 把它做成了可配置的层列表。
• main_kv_proj 是一个 ReplicatedLinear，一次把 hidden_size 投影成 MTP_num_layers * qk_head_dim。一次矩阵乘给所有草稿层产出 K，再用 unflatten(-1, (MTP_num_layers, qk_head_dim)) 拆成「每层一份」。这比逐层做一次投影省掉了 MTP_num_layers - 1 次 kernel 启动。
• _clone_linear_weight 与注释 main 和 draft 投影需要不同的连续布局，KV 权重会保留两份，显存会略微增加 说明 DFlash 的 KV 投影权重存了两份。原因是主模型路径与草稿路径对连续布局的要求不同，代码选择用一点显存换掉运行时的转置拷贝。

关于「什么时候写、写到哪个 slot」，docstring 里有段关键论证：

1
Decode 时随主模型 CUDA graph 一起 replay，此时接受结果尚未产生，全部 token 都按
2
slot_mapping_store_swa 写入：将被拒绝的 draft 位置正是 forward_drafts 随后写 draft key
3
的位置，在任何读取之前就会被覆盖，因此无需按接受结果屏蔽。

这段话回答了一个很自然的疑问：主模型前向发生在验证之后，但它写 KV 时还不知道哪些草稿被接受，怎么办？ 答案是「不需要知道」——被拒绝的那些位置随后会被 forward_drafts 写入的草稿 key 覆盖掉，而覆盖发生在任何读取之前。这是一个用写入顺序而非条件判断来保证正确性的典型手法，它让 update_main_cache 可以直接挂在主模型图里、不必等验证结果回传。

一轮 decode 的完整时序

run_drafts 是草稿侧的「整轮」入口，它的 docstring 明确说了这段是被 CUDA Graph 作为整体捕获的：

1
With the DFlash CUDA graph this is captured as a unit: the driver's broadcast
2
(followers receive into buffers allocated here, whose addresses the graph fixes),
3
the fixed-stride metadata, the draft layers and the sampling (which includes the
4
pool write) are all replayed.

代码只有六步：

1
cp_broadcast_multi_tensor([last_valid_offset, golden_tokens])
2

3
dflash_input_metadata = main_input_metadata.build_dflash_input(
4
    decode_input.first_token_positions,
5
    last_valid_offset,
6
    self.dflash_block_size,
7
    self.window_size,
8
)
9
hidden_states_before_norm = self.forward_drafts(golden_tokens, dflash_input_metadata, nextn_kv_caches)
10
if driver_input is None:
11
    return None
12
return self.sample_drafts(hidden_states_before_norm, driver_input)

把一轮 decode 从「主模型产出」到「下一步输入」串起来，顺序是这样的（$N$ 表示第 $N$ 步）：

1. 第 $N$ 步主模型前向：输入是本轮要验证的 1 + nextn 个位置。前向里顺手执行 update_main_cache，把主模型特征写进草稿侧的 main cache。同时产出 hidden_states。
2. 第 $N$ 步验证（采样侧）：Sampler._verification 用主模型的 logits 对草稿做拒绝采样，得到 last_valid_offset（接受数）与 golden_tokens（本轮最终确定的 token）。关键：golden_tokens 不是草稿，而是「被接受的草稿 + 第一个被拒绝位置重采样的 token」——按论文的说法它是本轮真正的输出，也是下一步草稿生成的锚点（anchor）。
3. 第 $N$ 步草稿生成：mtp_module.decode_forward。以 golden_tokens 为起点，DFlash 并行吐出 dflash_block_size 个草稿 token 与它们的置信度，写进 MTPInfoPool。
4. 第 $N$ 步后处理：process_MTPInfo 把 draft token、置信度累乘后的 accept_rates 记进 C++ 侧的 MtpState，供第 $N+1$ 步的调度使用。
5. 第 $N+1$ 步调度：用第 $N-1$ 步之前的置信度决定本轮验证容量，用第 $N$ 步的置信度决定每个请求的 seqlen_q（详见下节）。

这里有一个容易被忽略的细节：cp_broadcast_multi_tensor 是所有 rank 都要参与的，但只有 driver rank 有 driver_input。follower rank 拿到的是 torch.empty_like(...) 占位张量，广播把它们填上。代码里的断言说明了前提：

1
assert PG.context_parallel_group().size() > 1, "DFlash MTP follower ranks require context parallelism"

也就是「只有在 CP > 1 时才存在 follower rank」——单机单组时 driver 就是唯一的 rank。这个设计让草稿阶段的权重矩阵乘可以在 EP/CP 组内共享，而不必每个 rank 各跑一遍草稿（虽然它同时也意味着草稿结果要经过一次广播同步）。

fixed-stride metadata：草稿注意力要什么样的元数据

DFlash 的草稿层做的是块内因果注意力：同一请求的 dflash_block_size 个草稿位置彼此可见，同时它们都还要看到 SWA 窗口内的历史 KV。这跟主模型自回归解码时「一 token 看全部历史」的形态完全不同，因此需要一份专用的注意力元数据。

build_dflash_input 就是干这个的，它在 input_metadata/dflash.py 里：

1
(
2
    out.dflash_newx_slot_mapping_store,
3
    out.positions,
4
    out.input_ids,
5
    out.valid_token_mask,
6
    out.gate_force_random,
7
    out.dflash_kv_index,
8
) = get_fused_dflash_get_input(
9
    self.slot_mapping_store_swa,
10
    self.positions,
11
    first_token_positions,
12
    last_valid_offset,
13
    swa_block_size,
14
    dflash_block_size,
15
    window_size,
16
    has_valid_token_mask=(self.valid_token_mask is not None),
17
    has_gate_force_random=(self.gate_force_random is not None),
18
    input_ids_dtype=self.input_ids.dtype,
19
)

两个输出字段的注释说明了一切：

1
dflash_newx_slot_mapping_store: torch.Tensor | None = None
2
"""swa_cache store slots of draft-token kv. Shape: [real_cnt * dflash_block_size]."""
3
dflash_kv_index: torch.Tensor | None = None
4
"""swa_cache slots of the keys each draft token attends to; -1 = invalid.
5
Shape: [real_cnt * dflash_block_size, pad64(window_size + dflash_block_size)]."""

这是整个 DFlash 里最有信息量的两行。 逐项读：

• dflash_newx_slot_mapping_store 长度是 real_cnt * dflash_block_size，即每个草稿位置都有自己一个写槽位——草稿 token 的 K 也要进 SWA 缓存。
• dflash_kv_index 的形状是 [real_cnt * dflash_block_size, pad64(window_size + dflash_block_size)]。它的第二维同时包含 SWA 窗口内的历史 KV 槽位 与 块内 dflash_block_size 个草稿槽位。这说明草稿的注意力是「历史窗口 + 块内」两部分拼起来的一次 gather，而且第二维被 pad 到 64 的倍数——这是一个非常明确的 kernel 对齐要求（FlashMLA 系 kernel 对 seqlen_k 有对齐偏好）。
• -1 = invalid 是掩码约定。哪些位置会是无效的？块内超过「已接受长度」的草稿位置。这是下一段要讲的变长验证在注意力侧的直接体现。

为什么必须是「fixed-stride」？因为每个请求的块宽固定为 dflash_block_size，不随接受长度变化。这样整个张量的形状在编译期就确定，CUDA Graph 才能捕获它。变长只体现在 -1 掩码的分布上，不体现在张量形状上——这是把变长语义塞进静态图的标准做法。

DFlashMetadataMixin 的注释把「谁生产、谁消费」也写清楚了：

1
字段不由 C++ prepare 产出：decode 时 build_dflash_input 在 MTP 阶段构造（用 DFlash
2
cudagraph 时随图捕获、随图 replay）；消费者是 DFlash draft attention（attn.py）。

这些字段不走 C++ 的 prepare 路径，而在 Python 侧的 MTP 阶段产生。这跟主模型的注意力元数据（由 csrc/ 里的 C++ 代码准备）形成了对比，原因是草稿元数据依赖 last_valid_offset 这个「验证结果」，而验证结果是采样的产物，不在调度器的知识范围内。

DFlash 自己的 CUDA Graph

如果草稿阶段每步都逐 kernel 启动，那么「省下来的解码步数」会被启动开销吃掉一部分。DFlash 因此有独立的 CUDA Graph Runner（worker/dflash_cuda_graph_runner.py），与主模型的图分开放。

它的捕获循环：

1
self.batch_sizes = build_graph_batch_sizes(max_num_seqs, envs.DEEPINFER_DFLASH_CUDAGRAPH_BATCH_SIZE_INTERVAL)
2
...
3
for num_seqs in reversed(self.batch_sizes):
4
    decode_input = decode_input_full.cut(num_seqs)
5
    output = output_full.narrow(num_seqs) if self.do_sample else None
6

7
    # warmup
8
    graph_stream.wait_stream(torch.cuda.current_stream())
9
    with torch.cuda.stream(graph_stream):
10
        module.run_drafts(main_input_metadata, decode_input, nextn_kv_caches)
11

12
    graph = torch.cuda.CUDAGraph()
13
    with torch.cuda.graph(graph, pool=graph_memory_pool, stream=graph_stream, capture_error_mode="thread_local"):
14
        dflash_output = module.run_drafts(main_input_metadata, decode_input, nextn_kv_caches)
15
        if self.do_sample:
16
            output.copy_(dflash_output)
17
        del dflash_output

几个设计决策及理由：

其一，按 batch size 分桶捕获，从大到小。 reversed(self.batch_sizes) 从最大的桶开始捕获。原因在第四篇讲过同类问题时有提及，这里是显存：CUDA Graph 的捕获需要一块图形内存池（graph_memory_pool，由调用方传入），先捕获大桶能让这块池按最大需求分配一次，后续小桶复用同一个池。反过来从小到大捕获，池会反复增长，最终还是会到大桶的量级，只是多花了扩容的开销。

其二，输入输出都预先分配到 max_num_seqs，再用 .cut() / .narrow() 取前缀视图。

1
self.decode_input_padding = MTPDecodeInput.build_padding(self.max_num_seqs, module.noise_token_id, self.do_sample)
2
decode_input_full = MTPDecodeInput.build_padding(self.max_num_seqs, module.noise_token_id, self.do_sample)
3
output_full = module.allocate_output(self.max_num_seqs) if self.do_sample else None

decode_input_full.cut(num_seqs) 拿到的是同一块缓冲的前缀视图，output_full.narrow(num_seqs) 同理。这样图里引用的地址在小桶复用时仍然有效——因为小桶的输入就是大桶缓冲的开头部分。这是一个「零拷贝切桶」的技巧，与前面 get_or_create_stacked 注释里提到的 q_stride 可能大于 size[0] 是同一类操作。

其三，del dflash_output 是有意的。

1
dflash_output = module.run_drafts(...)
2
if self.do_sample:
3
    output.copy_(dflash_output)
4
del dflash_output

output.copy_(dflash_output) 让 output 承担了「图外的输出缓冲」角色，而 dflash_output 是图内临时张量。显式 del 是为了在图捕获结束时就释放它——否则这个张量会被 Python 引用着，占住图形内存池的一部分。捕获期的显存水位直接决定后续能不能跑起来，所以这段清理不是洁癖。

其四，选择用小桶向上取整。

1
def select(self, req_cnt: int) -> tuple[int, DFlashCudaGraphData]:
2
    index = bisect_left(self.batch_sizes, req_cnt)
3
    if index == len(self.batch_sizes):
4
        raise ValueError(f"No DFlash CUDA graph for {req_cnt} sequences; maximum captured batch size is {self.batch_sizes[-1]}")
5
    num_seqs = self.batch_sizes[index]
6
    return num_seqs, self.graphs[num_seqs]

bisect_left 找到第一个不小于 req_cnt 的桶，用这个桶的图跑。多出来的尾部位置由 decode_input_padding 填——它保存的是噪声 token（noise_token_id）填充的空请求。replay 时：

1
def replay(self, decode_input: MTPDecodeInput) -> MTPPredictInfo | None:
2
    req_cnt = decode_input.req_cnt
3
    _, data = self.select(req_cnt)
4
    copy_dflash_decode_input(data.decode_input, decode_input, self.decode_input_padding)
5

6
    data.graph.replay()
7
    if self.do_sample:
8
        return data.output.narrow(req_cnt)  # remove CUDA graph padding

copy_dflash_decode_input 把真实输入拷进图内缓冲、不足的用 padding 补齐；replay 之后 output.narrow(req_cnt) 把 padding 部分裁掉。「向上取整 + 填充 + 裁剪」是 CUDA Graph 处理动态 batch 的标准三段式，本站拆 SGLang 与 vLLM 的 CUDA Graph 时看到的也是同一套。区别在于这里的 padding 用噪声 token 而不是重复真实 token——因为草稿阶段对 padding 位置算出来的东西会被丢掉，只需要它们不产生非法访存即可。

markov：给并行草稿补上块内依赖

回到论文的第一个问题——并行草稿的后缀衰减。DFlash 的答案是 markov_rank：

1
self.markov_rank = cfg.dflash_cfg.get("dspark_markov_rank", 0)
2
self.enable_markov = self.markov_rank > 0
3
...
4
if self.enable_markov:
5
    self.markov_w1 = VocabParallelEmbedding(cfg.vocab_size, self.markov_rank, weight=...)
6
    self.markov_w2 = ReplicatedLinear(self.markov_rank, cfg.vocab_size, use_fp32_output=True, weight=...)

在采样草稿时：

1
for step_i in range(self.max_MTP_value):
2
    logits = logits_next[step_i :: self.max_MTP_value]
3
    if self.enable_markov:
4
        w1 = self.markov_w1(draft_token_ids)
5
        logits = self.markov_w2(w1, out_acc=logits)

逐项读：

• markov_w1 是一个词表嵌入，把「上一个草稿 token」映射到 markov_rank 维的低秩空间；markov_w2 再映射回词表维度并累加到当前步的 logits 上（out_acc=logits）。
• 输入的 draft_token_ids 在循环里被不断更新——每一轮 process_sampling 会把样本写回 draft_token_ids。所以第 $k$ 步的 logits 被第 $k-1$ 步采出的草稿 token 修正。这就是论文所说的 coupling a parallel backbone with a lightweight sequential module —— 骨架（并行主干）给出初始 logits，顺序模块（这里的低秩 Markov 项）按步注入前一个 token 的影响。
• 为什么是低秩？ 因为顺序修正必须在并行前向之后、逐步串行执行，它是关键路径上的额外开销。markov_rank 很小（配置文件给出，量级是个位数到几十），w1 是一次 embedding gather、w2 是一次小矩阵乘，相比重跑一遍草稿层便宜得多。这是「花小钱补上大依赖」的取舍。
• use_fp32_output=True 是精度要求：logits 的累加必须在 fp32 里做，否则在 bf16 下累加会丢有效位——尤其 markov_w2 输出的是修正量，量级可能远小于主 logits，用低精度累加会被吞掉。

论文把这一点放在「draft quality」这条线下，与它在离线 benchmark 上报告的「接受长度大幅超过 SOTA 的自回归与并行草稿器」相呼应。代码里 MTP_num_layers（草稿层数）与 markov_rank（顺序修正维度）是两个独立的设计旋钮：前者是并行骨架的深度，后者是顺序依赖的表达能力。

prefill_only：草稿头也可以只做一半

1
self.prefill_only = cfg.engine_type == EngineType.PROMPT
2
if self.prefill_only:
3
    logger.info("DFlash prefill_only 模式：只加载 main-cache 路径参数")
4
...
5
if not self.prefill_only:
6
    self.post_norm = ...
7
    # markov / hc_head / confidence_head 都不建

在 P/D 分离部署里，prefill 实例不生成草稿（prefill 只做 prompt 前向，一次就产出几千 token，投机没有收益），但它仍然需要 update_main_cache——因为主模型特征要写进 KV，而这份 KV 会随 P/D 传输交到 decode 实例手上。所以 PROMPT 引擎里的 DFlash 只保留 main-cache 路径参数，丢掉采样侧的全部结构（post_norm、markov_w1/w2、confidence_head）。这是一个很干净的做法：用同一个配置开关把「结构」和「权重」一起裁掉，而不是在运行时判断。

ConfidenceHead 的形状

1
class ConfidenceHead(nn.Module):
2
    """Predict the acceptance probability of each draft token."""
3

4
    def __init__(self, n_embed, markov_rank, num_steps, weight):
5
        self.proj = ReplicatedLinear(
6
            input_size=n_embed + markov_rank,
7
            output_size=1,
8
            weight=weight,
9
        )
10
        temperature = envs.DEEPINFER_DFLASH_CONFIDENCE_HEAD_TEMPERATURE
11
        temperatures = temperature if isinstance(temperature, tuple) else (temperature,) * num_steps
12
        self.temperature = tuple(max(value, SAMPLING_EPS) for value in temperatures)

三个设计点：

其一，输入是 n_embed + markov_rank。 也就是 hidden state 拼接上 markov 低秩嵌入。为什么要把 markov 特征也喂进去？因为 c_k 建模的是条件概率「在前面都通过的条件下，第 $k$ 个位置通过」。前面的 token 已经通过这件事的信息，正好浓缩在 markov 嵌入里。不接这个输入，头就只能猜无条件概率，而这正是 STS 要修的那种过度自信。

其二，输出维度是 1。 每个草稿位置一个标量，不做词表级的预测。这跟第四篇讲的 ParallelRoutedLinear 那类「输出整个词表」的结构完全不同：置信度头是个极小的模块，开销可以忽略。

其三，权重在装载时被显式转成 bf16：

1
weight=modules[-1].confidence_head.proj.weight.to(torch.bfloat16),

这一行说明模型文件里的置信度头权重可能是更高精度（fp32）。转 bf16 之后再配合 confidence_head_fused kernel 使用，是为了让整个头跑在一次融合 kernel 里。代价是置信度本身的精度降低——但这可以接受，因为它只是一个调度启发式信号，真正的正确性由拒绝采样保证。这是一个很值得记住的分工原则：验证保证无损，置信度只负责效率，所以两者可以有不同的精度标准。

融合 kernel 与「按步不同的温度」

1
def forward(self, hidden, markov_emb, out=None, step_i=0):
2
    return confidence_head_fused(
3
        hidden,
4
        weight=self.proj.weight,
5
        temperature=self.temperature[step_i],
6
        markov_emb=markov_emb,
7
        out=out,
8
    )

confidence_head_fused 把「拼接 → 线性投影 → 温度缩放 → sigmoid」压成一个 kernel，并且支持 out= 原地写入。调用侧：

1
if confidence_rates is not None:
2
    self.confidence_head(
3
        hidden_states_before_norm[step_i :: self.max_MTP_value],
4
        w1,
5
        out=confidence_rates[:, step_i],
6
        step_i=step_i,
7
    )

三点：

• hidden_states_before_norm[step_i :: self.max_MTP_value] 是跨步切片：所有请求在第 step_i 个草稿位置的行。这个步长切片（stride slice）是「按位置分组」而不是「按请求分组」，因为温度是按位置配的。
• out=confidence_rates[:, step_i] 直接把结果写进输出张量的一列，避免中间分配。
• step_i 只用来选温度 self.temperature[step_i]。这正是 STS 的落地形态：每一位有自己的温度，校准后写进环境变量元组即可，代码无需改动。

c_k 从设备到 CPU 的路径：累乘发生在 C++

置信度算完之后要跑一段不短的路才能变成调度信号。第一站是 sample_drafts 的返回：

1
prediction = MTPPredictInfo(token_ids=next_n_token_ids, confidence_rates=confidence_rates)
2
self.mtp_info_pool.store(
3
    sampling_input.mtp_info_write_slots,
4
    prediction,
5
    next_n_token_logprobs,
6
    next_n_token_logprobs_index,
7
)
8
return prediction

MTPInfoPool.store 把「草稿 token / 每个位置的 logprobs / 置信度」按 scheduler slot 写进一块稳定的 GPU 存储。池子的注释把显存账算得很清楚：

1
Each slot owns ``nextn`` rows. Slot 0 is reserved for padding; all other
2
slots store one row per MTP proposal.
3

4
``pool_rows = (max_num_seqs + 1) * nextn``. Logprobs occupy
5
``pool_rows * (force_topk or vocab_size) * 4`` bytes; force-top-k mode
6
allocates the same amount again for int32 indices.

为什么用「按 slot 索引」而不是「按 batch 位置索引」？ 因为草稿是在第 $N$ 步产出的，而它要被第 $N+1$ 步的验证读取——这两步之间请求集合可能变化（有新请求进来、有请求结束）。用 scheduler 分配的 info_slot 作为索引，跨步的对应关系就稳定了，不会因为 batch 重排而错位。这跟第四篇 EPLB 里 expert_count 按 physical 槽位计数是同一个思路：用一个稳定索引把跨步的张量绑在一起。

logprobs 在 force_topk 模式下只存 top-k，logprobs_index 存对应的词表编号。force_topk == 0 时存整个词表——显存差别是几个数量级，所以 DEEPINFER_SAMPLER_FORCE_TOPK 是个实打实的显存开关。注释里那两行初始化也值得一提：

1
# The scheduler uses slot 0 before a real slot has produced proposals.
2
self.logprobs[: self.slot_width] = float("-inf")
3
self.logprobs[: self.slot_width, 0 if force_topk else self.pad_token_id] = 0

slot 0 是 padding 槽，它的 logprob 分布被构造成「在 pad token 上概率为 1」。这样「还没产出草稿」的请求如果被误读，读出的是一个明确的哨兵值而不是垃圾数据——它会在拒绝采样时被判成负 logprob，从而必然被拒绝。用良构的哨兵值代替未定义行为是生产引擎的常见做法。

第二站：process_MTPInfo（C++ 侧）。这段代码是整条链路上最值得细读的一段：

1
mtp.prev_accept_rates = std::move(mtp.accept_rates);
2
mtp.accept_rates.clear();
3
mtp.prev_confidence_raw = std::move(mtp.confidence_raw);
4
if (has_confidence) {
5
    float accept_rate = 1.0f;
6
    for (int j = 0; j < nextn; j++) {
7
        accept_rate *= confidence_data[i][j];
8
        /// @note modeling返回的confidence rate是每个token独立的， 使用前要做一个累乘
9
        mtp.accept_rates.push_back(accept_rate);
10
        mtp.confidence_raw.push_back(confidence_data[i][j]);
11
    }
12
}

这段就是论文里 $a_k = \prod_{i \le k} c_i$ 的实现。 要注意它保留了两份：

• confidence_raw：原始的条件概率 $c_k$，注释说它「仅供 metrics 读取」。
• accept_rates：累乘后的前缀存活概率 $a_k$，供调度使用。

为什么要留 raw？因为 STS 校准的输入/输出关系需要原始值才能核对——如果只留累乘值，事后想诊断「校准是否生效」就没法把 $c_k$ 与 $a_k$ 分开看。同时 prev_confidence_raw 保留「上上轮的原始置信度」是给指标用的：比较两轮的置信度分布可以判断流量特征是否漂移。这与第四篇 EPLB 里保留 logical_balance 与 rank_balance 两个量、分别描述「问题输入」与「解决结果」是同一个观测哲学。

还有一处容易被略过的顺序约束：

1
/// @note build output 会读取更新前的 confidence_raw（即上面移入 prev_confidence_raw 的值），
2
/// 因此 processMTP 必须在 build output 之前调用

process_MTPInfo 必须在构造输出对象之前调用，因为 build output 读的是移动之前的值。这种「先记账后出账」的顺序要求在 llm_engine.py 的注释里也出现过一次（对 accepted_tokens 的说明），说明 DeepInfer 的这一步是有意设计成带副作用的两阶段，读者必须按调用顺序读才不出错。

「滞后」是怎么产生的：两轮 vs 三轮

论文里那句「用两步之前的置信度决定容量」听起来抽象，代码里的实现是双队列：

1
// 需要保证 confidence_{i+1, ..., nextn} 不影响 draft_token_i 是否被验证。
2
// 假设此时在第 N 轮，我们期望使用 N-2 轮的 next_n_token_accept_rate 去计算 best_bsz (total_q)；
3
// 并使用 N-1 轮的 next_n_token_accept_rate 去分配每个请求的 seqlen_q (non-ZOS 在这里做，ZOS 在 GPU 上修正)。

拆开这段话：

• 容量 $K$（best_bsz / total_q）用第 $N-2$ 轮的 accept_rates 算。为什么可以这么旧？因为容量只需要一个标量估计，流量特征是慢变量，两轮之前的统计仍然有代表性。
• 每个请求的 seqlen_q 用第 $N-1$ 轮的 accept_rates 分配——更新鲜，因为这是逐请求的决策，个体差异比全局总量变化快。

判定「哪个值是哪一轮的」靠两个计数器：

1
// 对于每个请求，维护 schedule_cnt 和 process_cnt 两个 counter，只有以下 2 种情况：
2
// 1. schedule_cnt == process_cnt，说明请求每次调度上 engine 后都产生了 accept_rate：
3
//    accept_rates 是 N-1 轮的值，prev_accept_rates 是 N-2 轮的值。
4
// 2. schedule_cnt == process_cnt + 1，说明请求最新一次 forward 还没产生 accept_rate（如 ZOS）：
5
//    accept_rates 是 N-2 轮的值，prev_accept_rates 是 N-3 轮的值。

这个「两个计数器 + 一次移位」的结构，本质是给每个请求维护一段两步长的置信度历史，然后在调度时按当前进度取对应的一轮。它也解释了 process_MTPInfo 里那句 mtp.prev_accept_rates = std::move(mtp.accept_rates) —— 移位就发生在这里。

那个断言则把这套账的边界钉死了：

1
TORCH_CHECK(mtp.process_cnt <= mtp.schedule_cnt && mtp.schedule_cnt <= mtp.process_cnt + 1, ...);

差值只允许是 0 或 1。 超过 1 说明有一步的记账丢失或重复——在 ZOS 的异步路径上这是很容易发生的事故（被抢占的请求计数会整组复位，process_MTPInfo 里专门有一段跳过逻辑处理它）。用一条断言把一个异步不变式写进代码，比让它悄悄产生错误调度要好得多。

与论文的对应关系

把代码与论文并排看，对应关系非常整齐：

变长投机把 decode 步变成一个容量分配问题

普通 decode 每步每请求 1 个 token，所以「本步总共算多少 token」＝活跃请求数，没有选择余地。开了变长投机之后，这个等式被打破了：每请求可以验证 $1 \sim 1+\gamma$ 个位置，于是本步的 total_q 成了一个需要决策的量。

回顾第二篇讲的 ZOS：下一步的 batch size 必须在当前步完成之前就已知，否则 GPU 流水线会 stall。所以这个「容量决策」不能等到验证结果出来再做——这正是 DSpark 必须引入置信度预测的根本工程原因（论文原话：Because ZOS requires the batch size for the next step to be known before the current step completes, synchronous scheduling would inevitably stall the GPU pipeline）。置信度头不只是效率工具，它是 ZOS 能在这个引擎里工作的前提之一。

MTPVarlen::schedule 的完整算法

C++ 侧的实现（csrc/engine/mtp_varlen.hpp）不长，但每一步都有讲究。先看它的「无曲线」退化分支：

1
if (_decode_perf_curve.empty()) {
2
    for (const auto &seq : scheduled) {
3
        seqlens_q.push_back(1 + static_cast<int>(seq->mtp.draft_tokens.size()));
4
    }
5
    return;
6
}

没有性能曲线时，退化成「固定长度验证」——每个请求验证它的全部草稿。这是一个「配置未提供则行为退化到基线」的设计，而不是报错或随便猜。它让这个特性可以在没有先做性能 profile 的环境里也能跑（只是没有加速收益）。

有曲线时的主流程分四步。

第一步，收集每 (请求, 位置) 的置信度并排序。

1
auto push_rates = [](std::vector<std::pair<float, int>> &queue, const std::vector<float> &rates, int nd, int seq_idx, bool has_draft) {
2
    if (rates.empty()) {
3
        // 若有 draft，pad 1.0，否则 pad 0.0
4
        queue.insert(queue.end(), nd, std::make_pair(has_draft ? 1.0f : 0.0f, seq_idx));
5
        return;
6
    }
7
    TORCH_CHECK(static_cast<int>(rates.size()) == nd);
8
    for (int j = 0; j < nd; ++j) {
9
        queue.emplace_back(rates[j], seq_idx);
10
    }
11
};

rates 为空时用 1.0 或 0.0 填充，而不是跳过。这个选择有实际含义：

• 有草稿但还没有置信度（例如刚做完 warmup、或者 ZOS 路径上置信度还在 GPU 上没回来）→ 填 1.0。语义是「乐观地假设它会通过」，于是这个请求的草稿位置会排在队列前面，优先获得验证预算。
• 没有草稿（nd == 0 且 has_draft 为假）→ 填 0.0。它根本不参与竞争。

「缺信息时取乐观值」在这里是合理的，因为漏掉一个本来能通过的草稿，损失是实打实的吞吐；而多验证一个会被拒绝的草稿，损失只是一点算力。两个方向的代价不对称，所以默认偏向乐观。

排序函数：

1
auto by_rate_desc = [](const std::pair<float, int> &a, const std::pair<float, int> &b) {
2
    return a.first != b.first ? a.first > b.first : a.second < b.second;
3
};

按置信度降序；平局按请求序号升序。 平局规则与第四篇 EPLB 里 argmin 在同分时返回最小 rank ID 是同一个考量：确定性。置信度是 fp32，真实流量下平局极其罕见；但一旦出现，不同的 tie-break 会让不同 rank 算出不同的调度，而在 EP/CP 组内这会直接导致集合通信错配。用「序号小者优先」这个与 rank 无关的规则，保证所有 rank 一致。

第二步，在性能曲线上搜索最佳 batch size。

1
auto sampler = _sampler_time_by_num_seqs.find(num_seqs);
2
// 对于 sampler，如果相同 num_seqs 已经收集足够样本，就用相同 num_seqs 过往几步的时间估计
3
// 否则回退到用全局过往耗时估计
4
auto estimate = sampler != _sampler_time_by_num_seqs.end() ? sampler->second.get() : std::nullopt;
5
if (!estimate) {
6
    estimate = _sampler_time_fallback.get();
7
}
8
const double sampler_ms = estimate.value_or(0.0);
9
const double slope_ms = std::max(0.0, _forward_time_slope.get());
10

11
double best_throughput = 0.0;
12
int best_bsz = 0;
13
float cur_expected = static_cast<float>(num_seqs);
14
int cur_bsz = num_seqs;
15

16
auto update_best = [&]() {
17
    while (cur_bsz_idx < num_bsz_selections && cur_bsz > _decode_perf_curve[cur_bsz_idx].first) {
18
        ++cur_bsz_idx;
19
    }
20
    if (cur_bsz_idx == num_bsz_selections) return;
21
    const auto &[graph_size, baseline_ms] = _decode_perf_curve[cur_bsz_idx];
22
    const double time_ms = baseline_ms + slope_ms * graph_size + sampler_ms;
23
    const double cur_throughput = cur_expected / time_ms;
24
    if (cur_throughput > best_throughput) {
25
        best_throughput = cur_throughput;
26
        best_bsz = cur_bsz;
27
    }
28
};
29

30
update_best();  // All the golden tokens
31

32
for (auto &acc : acc_queue_for_bsz) {
33
    cur_expected += acc.first;
34
    cur_bsz += 1;
35
    update_best();
36
}

这段就是论文 Algorithm 1 的落地。逐项拆：

• cur_expected 从 num_seqs 起步。这 $R$ 个 token 是「每个请求必然要算的那 1 个 golden token」——也就是论文里的 anchor。它们永远不会被裁掉，因为无论投机与否，解码都要前进至少一步。
• 每多纳入一个候选草稿位置，cur_expected 加的是它的 $a_k$，而不是加 1。 这就是论文的期望接受长度 $\tau$ 在增量上的体现：一个位置的边际收益是它自己的存活概率。这一步是「按边际收益分配预算」的精确表述。
• 时间模型是 baseline + slope * graph_size + sampler。三项分别是：启动曲线给出的基线（由 _decode_perf_curve 提供）、按 graph size 线性修正的项、以及采样阶段耗时估计。
• slope_ms 被 std::max(0.0, ...) 夹到非负，注释解释了这个下界的作用：以启动曲线为下界，避免负斜率外推到更大的 graph 后产生非正耗时。用启动曲线当下界是一个安全性设计——曲线本身是低负载下测出来的，把它当基线保证了大 batch 的估计不会比小 batch 还快（那会导致除零或负吞吐）。
• cur_bsz_idx 单调向前推进而不回退：因为 _decode_perf_curve 按 batch size 严格递增，而 cur_bsz 也在递增，所以内层 while 的总复杂度是 $O(\text{桶数})$ 而不是每轮 $O(\text{桶数})$。整体复杂度 $O(R \log R)$（排序主导），而不是 $O(R \cdot \text{桶数})$。

那条「性能曲线」在真实硬件上长什么样？论文特意强调它不是光滑的单峰曲线，而是锯齿状的阶梯函数（原文：the true hardware capacity SPS(B) is inherently discrete, exhibiting a jagged, step-wise degradation）。原因很实在：不同 batch size 会落进不同的 CUDA Graph 桶、触发不同的 kernel 配置，于是「每秒能出多少 token」在相邻 batch size 之间可能出现台阶式跳变。这条锯齿曲线正是论文要删掉提前退出的动机——在一段局部平坦的区间里贪心搜索，可能停在台阶前，而全局搜索能跨过这道坎找到后面更高的平台。

官方部署文档里那张 Wide-EP 的 decode 吞吐对比图，可以看作这条曲线在真实集群上的一次实测：同一模型在不同 EP 规模下，吞吐随并发上升的形状完全不同——这也就是为什么调度器必须用实测曲线而非解析模型。

（图源：DeepSeek V4.1 Flash 推理部署 Q&A，见 DeepInfer 仓库 doc/images/wide_ep_decode_pareto.png）

第三步，把选中的容量摊到各请求：

1
TORCH_CHECK(best_bsz >= num_seqs);
2

3
seqlens_q.assign(num_seqs, 1);
4
for (int i = 0; i < best_bsz - num_seqs; ++i) {
5
    seqlens_q[acc_queue_for_seqs_q[i].second] += 1;
6
}

每个请求先保底 1 个位置，剩余的 best_bsz - num_seqs 个预算按第二个队列（用更新的第 $N-1$ 轮置信度排的序）从高到低发放。 因为队列已按降序排好，前若干个元素天然就是「全局置信度最高的那批 (请求, 位置) 对」——这正是论文说的 dynamic top-K selection。

注意这里有个很微妙的语义：seqlens_q[r] += 1 对同一个请求可能执行多次，所以一个请求的 seqlen_q 是「被选中的位置数 + 1」。但选中的位置必须构成前缀——如果按全局排序，某个请求可能有第 3 个位置入选而第 2 个没入选，这在语料上是无意义的（草稿必须连续前缀才能验证）。

代码靠两层机制避免这个问题：

• 在排序层面：同一个请求内部的置信度是累乘值 $a_k$，而 $a_k$ 是单调不增的（因为每个 $c_i \le 1$）。所以同一请求的位置在全局排序里天然按 $k$ 递增的顺序出现——选中的一定是前缀。这个性质不是巧合，而是「用累乘值而非原始值排序」的直接结果。如果用手写的 confidence_raw 排序，就会破坏前缀性质，这是一个很容易踩的坑。
• 在代码层面：acc_queue_for_seqs_q[i].second 是请求序号，多个位置选中同一请求时靠 += 1 累加长度。因为前缀性质成立，累加出来的长度就是合法的。

第四步，断言兜底。上面那条 TORCH_CHECK(best_bsz >= num_seqs) 保证了「保底 1 个位置」永远有预算。反过来说：如果性能模型认为「连 golden token 都不值得算」，那是模型算错了，而不是应该接受的结果。

为什么容量用旧值、逐请求长度用新值

把双队列放回论文的非预测性约束下看，分工就清楚了：

• 容量 $K$ 是全局标量。它决定「这一轮总共验证多少个位置」。论文说 the historical prediction from two steps prior is used solely to determine the dynamic truncation length (i.e., the batch capacity limit K)。因为 $K$ 只影响准入的规模，不影响「选谁」的原则——用旧值近似它，引入的只是一个轻微的时序偏移。
• 逐请求长度必须用最新的。论文：the candidate tokens in the current step are still strictly sorted by their actual, up-to-date cumulative confidence scores。排序用最新值、容量用旧值，两者合起来保证：最自信的草稿永远优先被验证，即排序保持（rank-preserving）。

如果反过来（容量用新值、排序用旧值），就会出现「用未来信息决定谁入选」的泄漏。这就是论文宁愿付出两轮时延也要这样分工的原因：把「不精确」限制在一个不影响正确性、也不影响优先级的维度上。

ZOS 路径：把分配搬到 GPU 上

非 ZOS 路径上，seqlens_q 在 CPU 侧算好后直接传给输入准备。ZOS 路径上就不行了——因为接受结果还没回来，CPU 不知道上一轮实际接受了多少。解决办法是把分配下沉到 GPU：

1
if envs.DEEPINFER_ENABLE_MTP_VARLEN:
2
    compute_seqlens_q_by_confidence_rates(total_q, remap[:n_seq], prev.sampler_output.confidence_rates_gpu, seqlens_q[:n_seq])
3
if self.nextn > 0:
4
    compute_seqlens_k_zos(
5
        seqlens_k[:n_seq], seqlens_q[:n_seq], remap[:n_seq],
6
        prev.sampler_output.num_accepted_tokens_gpu, nextn=self.nextn,
7
    )

compute_seqlens_q_by_confidence_rates 是 TileLang kernel，参考实现的语义如下：

1
# Phase 1: compute cumulative product of confidence rates via remap
2
mask = remap != -1
3
safe_remap = remap.clamp(min=0)
4
gathered_rates = prev_confidence_rates[safe_remap]
5
# For fresh sequences, set rates to 0.0 (so cumprod = 0.0 everywhere)
6
gathered_rates = torch.where(mask.unsqueeze(1), gathered_rates, torch.zeros_like(gathered_rates))
7
cum_rates = torch.cumprod(gathered_rates, dim=1)
8

9
# Phase 2: rank-based selection
10
extra_budget = total_q - num_seqs
11
flat_rates = cum_rates.reshape(-1)
12
sorted_indices = torch.argsort(flat_rates, descending=True, stable=True)
13
ranks = torch.empty(len(flat_rates), device=remap.device, dtype=torch.int32)
14
ranks[sorted_indices] = torch.arange(len(flat_rates), device=remap.device, dtype=torch.int32)
15
ranks = ranks.reshape(num_seqs, nextn)
16
wins = (ranks < extra_budget).to(torch.int32).sum(dim=1)
17
seqlens_q.copy_(1 + wins)

逐点对照：

• remap 把「上一步的 batch 索引」映射到「本步的请求位置」，新请求用 -1 表示。这是异步路径下必须的：上一步的 batch 与本步的 batch 不一定同序、也可能有增删。
• cumprod 就是 accept_rates 的 GPU 版。因为 ZOS 下置信度还在显存里没回 CPU，累乘必须在 GPU 做。
• 新请求（remap == -1）的 rates 置 0，于是 cumprod 全程为 0，ranks 必然很大，永远不会入选 extra_budget。 语义是「没有可验证草稿的请求不参与竞争」，与 CPU 路径的 pad 0.0 一致。
• Phase 2 是同一套 top-K：用 stable argsort 得到排名，ranks < extra_budget 的位置就是入选的，每请求 1 + wins。
• assert num_seqs * nextn * 4 <= 32 * 1024 只对 NPU 生效，注释说明原因是 UB（Unified Buffer）的容量上限 32KB——这是 Ascend 侧的片上缓冲约束，与第一篇讲的「双平台」问题直接相关。

ZOS 路径还有一个「把修正结果拷回来」的动作，注释解释了必要性：

1
if envs.DEEPINFER_ENABLE_MTP_VARLEN:
2
    # H2D, correction, and D2H are ordered on the current stream; only varlen rewrites this pinned source.
3
    seqlens_q_cpu[:n_seq].copy_(seqlens_q[:n_seq], non_blocking=True)
4
    d2h_evt = torch.cuda.Event()
5
    d2h_evt.record()

只有变长模式会重写这份 pinned 源，所以非变长模式下不需要这次设备到主机的拷回。这是第二篇讲的「把同步点从热路径移到观测路径」的同一个原则：拷贝是 non_blocking 的，用 event 标记完成，由消费方自己选择等待时机——关键路径上没有任何 .item() 或 synchronize()。

引擎侧的记账闭环

草稿与接受结果最终要回到调度器，让序列长度、KV 缓存与输出全部对齐。三个环节：

其一，接受结果以 flat 布局直通 C++（llm_engine.py）：

1
ftp = output.first_token_positions  # [req_cnt], CPU int32, or None
2
if ftp is not None and output.accepted_token_ids is not None:
3
    accepted_token_ids = output.accepted_token_ids  # [total_q] flat, 直接传给 C++
4
else:
5
    accepted_token_ids = torch.empty((0,), dtype=torch.int32)

注释写着 C++ 下游函数已支持 [total_q] flat 布局。这里省掉的是一次 reshape——验证的产物天然是「按 packed row」的，如果 C++ 侧要求 [req_cnt, nextn] 的二维布局，就要多做一次拷贝与重排。

其二，prompt 阶段要补 pad：

1
# FIXME decode 如果 input 没有 nextn, 则 pad 到 nextn
2
# NOTE 只有 normal engine 在这里 pad；PROMPT engine 不产生 nextn 信息，置空
3
if output.next_n_token_ids is None:
4
    if self.engine_type == EngineType.NORMAL:
5
        next_n_token_ids = self.get_nextn_pad_token_ids()
6
    else:
7
        next_n_token_ids = torch.empty((0,), dtype=torch.int32)
8
    next_n_token_ids = next_n_token_ids.tile((num_reqs, 1)) if envs.DEEPINFER_USE_NEXTN_DECODE else next_n_token_ids.unsqueeze(0)

get_nextn_pad_token_ids 返回 [vocab_size - 1] * nextn——用词表最后一个 id 当 pad。这与 MTPInfoPool 里 pad_token_id = vocab_size - 1 是同一个常量，两处必须一致。PROMPT 引擎置空是因为它根本不产生草稿（对应前面 prefill_only 的设计）。这里有个已标注的 FIXME：TODO(DFlash): pad token 的应该是 0——说明 pad 值的选择在团队内部还没有最终定论。

其三，运行时耗时回灌调度器（post_process）：

1
if step_state.forward_events is not None and step_state.sample_events is not None:
2
    self.scheduler.update_mtp_runtime(
3
        len(seq_metadata_list),
4
        step_state.graph_size,
5
        step_state.forward_events[0].elapsed_time(step_state.forward_events[1]),
6
        step_state.sample_events[0].elapsed_time(step_state.sample_events[1]),
7
    )

update_mtp_runtime 把实测的 forward 与 sampler 耗时喂回 ForwardRuntimeEstimator 与 SamplerRuntimeEstimator（即第二篇拆过的那两个估计器）。性能曲线不是静态配置，而是被实测持续修正的——这是整个变长调度能自适应的根本原因。触发条件也写得很谨慎（worker/base.py）：

1
record_runtime = bool(
2
    self.is_driver_worker
3
    and envs.DEEPINFER_ENABLE_MTP_VARLEN
4
    and self.nextn > 0
5
    and not is_prompt
6
    and not is_placeholder
7
    and graph_size > 0
8
)

五个条件：只有 driver rank 测（避免重复）、只在开变长时测、必须有草稿、不在 prompt 阶段测（prefill 的耗时特征与 decode 完全不同）、不测 placeholder（占位步没有真实计算）。这种「什么情况下不测」的枚举，比「什么情况下测」更能说明测量口径的复杂度。

基线是 MTP-1，而不是 MTP-5

论文有一个容易看漏但很关键的事实：线上长期使用的基线是 MTP-1（草稿长度 1），而不是更长草稿的 MTP-3/MTP-5。原因是论文原话——deploying a static multi-token drafter (e.g., MTP-3/5) strictly degrades aggregate throughput under high concurrency due to excessive verification overhead。

也就是说，长草稿在离线指标上更好看，但在高并发下反而拖垮总吞吐，所以线上宁可退回到「只投机 1 个 token」。这条事实本身就是本篇全部机制分析的前提：DSpark 要解决的不是「草稿够不够长」，而是「长草稿在高并发下能不能被安全地用起来」。MTP-1 被 DSpark 取代发生在 DeepSeek-V4-preview 发布两周之后。

离线指标：接受的草稿长度

论文用「平均接受长度 $\tau$」（每轮解码实际前进的 token 数，含目标模型产出的 bonus token）作离线指标，在四个目标模型（Qwen3-4B/8B/14B 与 Gemma4-12B）上对比三种草稿器：Eagle3（自回归）、DFlash（并行）、DSpark（半自回归）。训练口径统一（同一份数据、同一套目标层特征、草稿头层数 DSpark/DFlash 为 5、Eagle3 为 1，块长统一 7），且离线评测关掉了置信度调度器，强制所有草稿器提出固定长度的块——这样比出来的才纯粹是草稿质量。

结果是 DSpark 在全部模型与全部领域上领先：相对 Eagle3 的宏平均提升为 30.9%、26.7%、30.0%（4B/8B/14B），相对 DFlash 为 16.3%、18.4%、18.3%。markov_rank 那个低秩顺序模块的收益，就是这么兑现的。

论文还专门指出了一个领域效应：结构化任务的自然接受长度明显高于开放式聊天（以 Qwen3-4B 为例，数学与代码领域约 5.1~5.6，而聊天领域只有 3.5 左右）。这正是「固定验证长度必然浪费」的直接证据——同一套草稿在不同数据分布下的「合适深度」差了一倍。

置信度校准的实测效果

论文第 4.3.3 节给出了置信度头本身的诊断，其中「可靠性图」把过度自信问题画得很直观：校准前的预测分布整体偏向高置信区间，说明模型对草稿通过验证的概率估计偏高；后验的 STS 校准把曲线拉回对角线附近。论文报告校准前 ECE 在 3%8%、判别能力 ROC-AUC 在 0.810.90；STS 校准后平均 ECE 降到约 1%。

（图源：DSpark 图 6，Alpaca 数据集）

为什么这 3%~8% 的偏差必须修？ 因为在硬件感知调度器里，置信度的绝对值直接参与「期望接受长度 × 时间」的吞吐估算。判别能力只决定「谁排前面」，绝对幅度决定「值不值得验证」。过度自信会让调度器高估验证的收益，从而多验证本来不该验证的草稿——正好抵消了它要解决的问题。这是「排序正确还不够、必须绝对准」的一个明例。

一个反直觉的工程决策：删掉提前退出

论文第 5.2 节讲了一段在纯算法视角下会显得奇怪的改动：为了不被真实硬件的 SPS 曲线「锯齿」困在局部最优，DSpark 删掉了 Algorithm 1 里的 early-stopping break，改成不受约束的全局搜索。

这个改动看起来直接违反了前面讲的非预测性约束——搜索空间变大，似乎更容易泄漏未来 token。论文的解释是：异步调度本身构成了一道因果障碍。因为无约束搜索评估的是「两步之前的预测值」，准入决策与当前 token $x_{r,k}$ 的实际取值天然隔离，所以放大搜索范围并不会引入选择偏差。

这段论证值得单独记住，因为它体现了一个很成熟的工程判断：一个在同步语义下不成立的算法改动，在异步时序下可以被证明是安全的。反过来讲，如果 DSpark 没有先做「延后两步」这套异步适配，那么「全局搜索」这个能显著提升物理吞吐的改动就根本不能做。时序设计反过来打开了算法设计的空间——这是本篇里最值得带走的一条方法论。

变长 query 的物理执行

调度器决定出「每个请求验证多长」之后，还要把它变成可执行的 kernel 形态。论文第 5.3 节的做法与我们在 build_dflash_input 里看到的一致：把物理执行与逻辑序列追踪解耦——所有请求的所有 token 被压平（flatten）成同一个元素维度、按相同方式处理，而序列内部的依赖关系则通过一个标记张量（marker tensor）编码，交给稀疏注意力实现去解读。

论文还给出了改动面：在 DeepSeek-V4 结构上，只有 index-attention 与 compress 两个 kernel 需要改（对应第一篇提过的 index_attn.py 与 compress.py）。这个「只改两个 kernel」的结论，是「只服务一个模型」这个前提带来的红利——通用框架要支持变长 query，代价会大得多。

线上收益

线上对比是 DSpark-5（$\gamma = 5$）对 MTP-1，在 DeepSeek-V4-Flash 与 V4-Pro 的生产引擎上，用真实用户流量的遥测数据画吞吐-时延前沿。核心结论：在同等吞吐水平下，单用户生成速度提升 60%~85%；更重要的是，它在严格交互性约束下避免了吞吐的剧烈下滑，把原先不可达的性能档位变成了可达。

（图源：DSpark 图 8，DeepSeek-V4 生产引擎实采遥测。上半部分是不同并发下的总吞吐，下半部分是每请求的平均验证预算。「并发升、预算降」这条趋势，就是 MTPVarlen::schedule 在线上真实负载下被触发的结果。）