详解SGlang 最新DFlash投机解码算法的工程实现、overlap策略和改进点

大模型推理慢，一部分慢在自回归生成：

生成 token_1 -> 再生成 token_2 -> 再生成 token_3 -> ...
每一步都依赖前一步 token，GPU 很强，但一次只能往前走一个 token。投机解码的目标，就是想办法让模型一次往前多看几步。

DFlash可以理解为之前文章的Eagle和LLada结合体，它是一种diffusion LLM并行预测方案。它不再让 draft model 严格自回归地产生 token，而是给 draft model 一个固定长度的 mask block，让它一次性非因果预测多个token，然后再交给 target model 进行验证。

目录

1. DFlash算法解析

2. SGLang DFlash 的整体框架

3. SGlang DFlash decode阶段实现

4. DFlash Draft model KV注入实现

图里几个颜色可以这样对应：

• 蓝色Fused Target Context Feature：来自 target model 捕获的若干层 hidden states特征，作为重要的上下文信息输入到draft model。
  

• 黄色Target Decode Token：当前已被 target model 确认接收的 token，也就是下一轮 block 的第一个 token。
  

• 绿色Mask Token：draft model 要并行预测的位置。
  

Target Embedding -> [上一轮 target token, mask, mask, ...] -> Draft Layers -> Target LM Head -> draft tokens -> Target Model verify

第二条是 DFlash Draft mode KV materialization：Draft mode基于target context feature 融合target信息，从而提高预测精度的路径。

Target Model -> 捕获若干层 target hidden states -> 融合成 target context feature -> 写入 Draft KV Cache

所以 DFlash 的核心是用 target model 已确认的上下文特征增强 draft model提高预测精度，再让 draft model 做并行预测提高预测速度，完事儿，就是这么直接的两个优化。

--speculative-algorithm DFLASH--speculative-draft-model-path

其核心有两个 runner：

• Target runner: 目标大模型，负责 prefill、target verify、logits 计算，以及捕获若干 target layer hidden states。

• Draft runner: DFlash draft model，负责 mask block tokens的并行预测。

对应源码主要在：

python/sglang/srt/speculative/dflash_worker_v2.pypython/sglang/srt/speculative/dflash_info_v2.pypython/sglang/srt/models/dflash.pypython/sglang/srt/managers/overlap_utils.py

其中 DFlashWorkerV2.forward_batch_generation() 是主流程。它会根据当前 batch 是 prefill 还是 decode，分别执行：

prefill: target model prefill 捕获 target aux hidden materialize 到 draft KV cachedecode: 构造 [bonus, mask, ...] block draft model 并行预测 target model verify 计算 accept_len / bonus / commit_lens 把 accepted + bonus 对应 hidden 写入 draft KV cache 生成下一轮 next_draft_input

3. SGlang DFlash decode阶段实现
重点讲一下decode阶段的实现过程，假设 block_size = 10。DFlash 每轮构造一个固定长度 block：

[bonus, mask, mask, mask, ..., mask] #长度为 10

其中第一个 bonus 是上一轮 target model 确认过的 token，后面 9 个位置是 mask。draft model 会一次 forward，并行预测这些 mask 位置对应的 hidden states，再通过 target lm_head 转成 draft token：

[bonus, draft1, draft2, ..., draft9]

然后 target model 一次性 verify 整个 block，得到 target 的预测：

target_predict[0], target_predict[1], ..., target_predict[9]

接受规则即是需要连续匹配

draft1 == target_predict[0]draft2 == target_predict[1]draft3 == target_predict[2]...

一旦中间不匹配，后面的 draft token 全部拒绝。最终：

accept_len = 连续接受的 draft token 数bonus = target_predict[accept_len]commit_lens = accept_len + 1

这里 +1 是新的 bonus token。也就是说，即使一个 draft token 都没被接受，target model 也至少会提交一个 bonus token，使得投机解码能够继续下去

对应核心代码：

def compute_dflash_correct_drafts_and_bonus( *, candidates: torch.Tensor, target_predict: torch.Tensor,) -> Tuple[torch.Tensor, torch.Tensor]: ... matches = candidates[:, 1:] == target_predict[:, :-1] correct_len = matches.to(torch.int32).cumprod(dim=1).sum(dim=1) bonus = target_predict[torch.arange(bs, device=target_predict.device), correct_len] return correct_len, bonus.to(torch.int64)

可以用以下伪代码理解decode生成的逻辑：

block = [last_bonus] + [MASK] * (block_size - 1)draft_hidden = dflash_draft_model( input_embeds=target_embedding(block), draft_kv_cache=draft_kv_cache,)draft_tokens = target_lm_head(draft_hidden[1:])candidates = [last_bonus] + draft_tokenstarget_logits, target_aux_hidden = target_model.verify(candidates)accept_len, bonus = verify(candidates, target_logits)commit_lens = accept_len + 1materialize_target_hidden_to_draft_kv( target_aux_hidden[:commit_lens])next_round_input = bonus

4. DFlash Draft model KV materialization实现

本节基于target model捕捉hidden states以及hidden states转换为draft model的kv cache两部分来介绍Draft mode KV materialization。

4.1 基于 target model 捕获的若干层 hidden states，构造 DFlash draft model 的上下文

DFlash 的另一个关键点是draft model 并不是凭空预测。它需要 target model 捕获若干层 hidden states，把这些 hidden states 融合成上下文特征，最后再写入 draft KV cache，增强 draft model 生成 token 的准确性。

以 `LLaMA3.1-8B-Instruct-DFlash-UltraChat` draft model config 为例：

"dflash_config": { "mask_token_id": 128002, "target_layer_ids": [1, 8, 15, 22, 29]}

SGLang 会读取这些 target_layer_ids，然后让 target model 捕获对应层的 hidden states：

if self.spec_algorithm.is_dflash() and not self.is_draft_worker: ... dflash_draft_config = parse_dflash_draft_config( draft_hf_config=draft_model_config.hf_config ) ... self.dflash_use_aux_hidden_state = True self.dflash_target_layer_ids = dflash_draft_config.resolve_target_layer_ids( target_num_layers=int(target_num_layers), draft_num_layers=int(draft_num_layers), )

捕获到的 aux hidden 最后会通过 logits processor 放进 LogitsProcessorOutput.hidden_states，供 DFlash worker 使用。

if batch.forward_mode.is_extend() or batch.is_extend_in_batch:Target prefill: capture DFlash aux hidden states for prompt tokens.batch.capture_hidden_mode = CaptureHiddenMode.FULL batch_output = self.target_worker.forward_batch_generation(batch) ... self._append_target_hidden_to_draft_kv_by_loc( target_hidden=logits_output.hidden_states, cache_loc=batch.out_cache_loc, positions=positions, )

在 decode verify 后，SGLang 会把本轮 accepted + bonus 对应的 target aux hidden通过 commit_lens 写入 draft KV：

if new_seq_lens is None: new_seq_lens = prefix_lens + commit_lens.to(prefix_lens.dtype)...hidden = logits_output.hidden_states...self._append_target_hidden_to_draft_kv_by_loc( target_hidden=hidden.reshape(-1, hidden.shape[-1]), cache_loc=verify_out_cache_loc, cache_loc_2d=verify_out_cache_loc_2d, positions=positions, commit_lens=commit_lens,)

_append_target_hidden_to_draft_kv_by_loc() 的逻辑可以概括为：

target aux hidden -> draft_model.project_target_hidden() -> 每一层 draft attention 的 kv_proj_only() -> k_norm + rope -> 写入 draft token_to_kv_pool

对应源码：

ctx_hidden = self.draft_model.project_target_hidden(target_hidden)...for layer in self.draft_model.layers: attn = layer.self_attn layer_ctx_hidden = self.draft_model.prepare_context_hidden_for_kv( layer, ctx_hidden ) k, v = attn.kv_proj_only(layer_ctx_hidden) k = attn.apply_k_norm(k) k = attn.apply_k_rope(positions, k) ... self.draft_model_runner.token_to_kv_pool.set_kv_buffer_prefix_valid( attn.attn, cache_loc_2d, commit_lens, k, v, ... )

所以，通过DFlash Draft model融合，target model 中包含上下文信息的 aux hidden经过 DFlash draft model 投影和 K/V projection 得到 draft KV，保存到draft KV cache中供下一轮draft model使用

commit_lens = accept_len + 1new_seq_lens = prefix_lens + commit_lens

也就是说，下一轮从哪个位置开始，必须等本轮 target verify 完成后才知道。同时，下一轮还要知道两个东西：

1. 下一轮 block 从哪个 position 开始

1. KV headroom:提前给未来可能写入的 token 准备 KV slot 和 req_to_token 映射。
2. FutureMap: 在当前轮 forward 中途，把 new_seq_lens 和 bonus token 传给下一轮。

有了这俩机制后，Dflash也兼容scheduler和fwd的overlap了，接下来的两小节结合源码来具体分析这两个的具体操作

5.1 KV headroom
根据之前DFlash decode的介绍，大家应该有印象target model每轮最多可能写 block_size 个位置的token（因为base model每轮生成block_size个token），但有效提交长度要等 verify 后才知道。如果等 verify 后再分配 KV，CPU scheduler 和 GPU forward 就只能串行运行。

所以 SGLang 就引入KV headroom为未来 token 提前预留好 KV cache 空间，SGlang需要在DFlashDraftInputV2.prepare_for_decode()里确定需要预留的最大length空间reserved_len：

cur_alloc_len = int(req.kv_allocated_len)reserved_len = max(cur_alloc_len, committed_len + 2 * block_size)#committed_len:当前真实已提交的 token 长度。#cur_alloc_len:当前已经预分配 KV slot 的逻辑长度上界。

有了预留的reserved_len后通过req_to_token 将KV headroom从token位置映射到对应的物理 KV cache：

req_to_token[req_pool_idx, logical_position] = physical_kv_slot#req_pool_idx：batch_id#logical_position:token位置id#physical_kv_slot：token对应的kv cache

在本轮DFlash验证后得到new_seq_lens，CPU scheduler就可以通过req_to_token 开启下一轮 DFlash 的kv cache构造，为下一轮GPU forward做准备：

prefix_lens = new_seq_lenspositions = prefix_lens + [0, 1, ..., block_size - 1]cache_loc = req_to_token[req, positions]

总而言之也就是说，KV headroom 解决的是：下一轮kv cache 要写到哪里。

SGlang DFlash worker 在算出 new_seq_lens 后就通过FutureMap.publish保存本轮生成的最终length即new_seq_lens，这样下一轮即可知道下一轮真实起点。

if new_seq_lens is None: new_seq_lens = prefix_lens + commit_lens.to(prefix_lens.dtype)if on_publish is not None: on_publish(new_seq_lens)

CPU scheduler 把 on_publish 绑定为 FutureMap.publish，下一轮CPU scheduler即可以通过resolve获得new_seq_lens从而知道从哪里开始：

def resolve_seq_lens_cpu(self, batch: ScheduleBatch) -> None: ... if self.publish_ready is not None: self.publish_ready.wait() batch.seq_lens = self.new_seq_lens_buf[fi]

同理类似bonus token（即本轮验证得到的最后一个token）通过 FutureMap.stash() 存起来，下一轮CPU scheduler通过 _resolve_spec_extras() 再填回 draft_input.bonus_tokens。这样下一轮 DFlash worker 才能构造：

[bonus, mask, mask, ..., mask]

整体可以这样理解：

KV headroom: 先准备好未来 position 对应的 KV slot。publish(new_seq_lens): 告诉下一轮从哪个 position 开始。stash/resolve bonus: 告诉下一轮 block[0] 是哪个 token。

6. DFlash 性能分析

从公开结果看，DFlash 的性能表现亮眼。DFlash 论文中展示了基于 SGLang 框架的单并发评估结果，吞吐量可以提升 2.2x 到 4.8x，部分截图如下：

SGlang DFlash性能测试

工业界的案例说明了 DFlash 具备落地价值。6 月 9 日，小米 MiMo 团队联合TileRT 推出 MiMo-V2.5-Pro-UltraSpeed，引入并行预测的 DFlash，在一台标准 8 卡通用 GPU 节点上，把万亿参数大模型文本生成速度推到 1000 TPS，峰值可达 1200 TPS。

这些结果说明从学术研究到工业落地 DFlash 都展现出通过并行推理方向能够有效增加推理速度。但基于我个人对 SGLang 实现和 DFlash 算法机制的理解，我认为DFlash 仍然存在一些值得讨论和优化的部分。

7 DFlash 值得思考和优化的部分

7.1 基于 SGLang 的 DFlash 落地思考

首先是 overlap。DFlash overlap 和普通 SGLang overlap 的思想是一致的：减少 CPU scheduler 对 GPU forward 的阻塞，让 CPU 和 GPU 尽量流水化。

普通 SGLang overlap 只需要传递上一轮采样出来的 next_token，下一轮 seq_lens 基本就是 +1。但 DFlash 更复杂，因为每轮提交多少 token 取决于 target verify 的 accept_len。因此它需要额外维护以下变量：

new_seq_lensbonus tokencommit_lensdraft KV materialization 状态req_to_token 未来映射

这也意味着，DFlash overlap 的收益依赖一个前提：CPU 侧 workload 要足够大，且 CPU 调度确实是瓶颈。如果 GPU 本身已经很慢，overlap 能隐藏的 CPU 开销有限，DFlash 额外引入的 draft forward、target verify、KV materialization、headroom 管理反而可能吞掉收益。

比如模型很大或者 GPU 较弱时，在 RTX 6000D 这类卡上跑 27B/35B 模型，收益可能很差，甚至出现负收益。因为此时瓶颈主要在 GPU forward，而不是 CPU scheduling。

但在 B200 这类更强硬件上，Qwen3-8B 可以从约 11.4 ktok/s 提升到约 15.3 ktok/s，提升超过 33%。这说明在合适硬件和合适模型规模下，SGlang DFlash 的overlap才能发挥真正的作用。

其次是 SGLang 当前实现层面的限制。

第一，当前 DFlash 不支持 DP attention 和 pipeline parallel。原因是 DFlash 的状态非常依赖本地 request、KV pool、req_to_token 映射和 draft KV materialization。DP/PP 会让 hidden、KV、request ownership 分散到多个 rank 上，当前实现还没有完整处理这些跨 rank 状态同步。

第二，内存和 KV 管理复杂。DFlash 同时维护 target KV 和 draft KV，还会在 overlap 下 reserve：

committed_len + 2 * block_size

这提升了 overlap 能力，但也增加了显存压力。吞吐和延迟会受到 block_size、accept rate、headroom、page size、draft window size 的影响。

7.2 DFlash 算法本身的局限

除了SGlang工程实现的局限，DFlash 算法本身也有一些天然限制。

1.后缀衰减问题。DFlash 的所有 mask 位置是块级并行预测，块内 token 没有严格前后因果依赖。token间的预测互不影响，无法完全模拟自然文本的时序逻辑，导致 block 越靠后的 token 准确率越容易下降，大量后缀草稿会被 target model 拒绝。

2.固定长度验证带来的算力浪费。DFlash 每轮固定 verify 完整 block，无法提前识别低置信度尾部 token。高并发下，后缀低质量 token 大概率被拒绝，但仍然占用 target model verify 算力，造成 GPU 无效开销。

这也是为什么后续会出现类似 DeepSeek DSpark 这样的改进方向。简单来说，DSpark 是一种融合半自回归生成与置信度调度验证的推测解码框架，可以理解为针对 DFlash 两个核心问题做增强。

半自回归生成解决的是块内 token 无前后信息的问题。它在 DFlash draft model 无因果并行生成 draft token 的基础上，引入一个轻量级模块，例如 Sequential Markov / RNN Head，再做一次轻量级自回归修正，使后面的 token 能看到前面 token 的信息，从而缓解后缀准确率下降。

置信度调度验证解决的是固定长度验证浪费算力的问题。DSpark draft model 会先通过 Confidence Head 预测 token 接受率，然后系统侧依据当前负载动态调整验证长度。当轻负载时，GPU 资源相对充足，可以验证更多 token。
重负载时则优先验证高置信度、高接受率 token，减少低接受率尾部 token 的无效 verify。

这样一来，DSpark 则是在DFlash基础上进一步补足块内依赖和验证调度的问题。从而进一步追求预测得更准、验证得更聪明。

这套机制能让 DFlash 在合适硬件和合适 workload 上获得非常可观的收益。但它不是无条件加速。它对 CPU/GPU 性能比例、模型大小、accept rate、block_size、KV 管理都有明显依赖。

所以我更倾向于把 DFlash 看成一种“强工程耦合型”的推理加速方案：算法本身提供并行预测能力，SGLang runtime 负责把这部分能力与runtime本身耦合兼容，从而转化成真实吞吐收益。至于真正能不能赚到性能收益，以及能赚到多少，还是要看模型、硬件和调度开销是否匹配，你的workload瓶颈在哪里。