DFlash:块扩散 + KV 注入,让投机解码的草稿不再串行
来源: The next generation of speculative decoding: DFlash and Spec V2 — Z Lab / Modal / SGLang,2026-06-15
模型: z-lab/Qwen3.5-397B-A17B-DFlash · modal-labs/Qwen3.5-397B-A17B-DFlash · lmsys/Qwen3.5-397B-A17B-DFlash
前置阅读: Speculative Decoding · EAGLE 系列 · MTP
一句话总结: EAGLE / MTP 把草稿模型砍到只剩一两层,但草稿本身仍然是逐 token 自回归的,对 GPU 不友好。DFlash 用「块扩散」一次并行吐出整块草稿、用「KV 注入」把 target 的上下文特征直接灌进草稿模型的 KV cache,同时压低草稿开销、抬高接受率,在 Qwen3.5-397B-A17B 上相对 baseline 拿到 >4.3×、相对原生 MTP 拿到 1.5× 的吞吐。
一、瓶颈:草稿还是串行的
经典投机解码的逻辑是「小模型猜、大模型并行验」。后来的 EAGLE 系列、MTP,以及 Gemma 4、DeepSeek-V4 等模型内置的多 token 预测模块,都把草稿模型做得越来越小——EAGLE 甚至只剩一层 Decoder。
但有一件事一直没变:草稿 token 仍然是一个一个自回归生成的,只不过自回归的位置从 target 模型挪到了草稿模型上。
1 | EAGLE / MTP 的草稿阶段: |
这正是 decode 阶段的老毛病:逐 token 自回归是 memory-bound、低算术强度的操作,跟 GPU/TPU 喜欢的大块并行计算南辕北辙。为了把草稿开销压住,这类方法只能用极浅的草稿模型(1~2 层),而浅模型又限制了草稿质量。于是陷入两难:
1 | 草稿做深 → 接受率高,但每步串行开销大 → 端到端不划算 |
DFlash 想同时打破这两端的约束:草稿生成不再串行,草稿模型也不必为了省钱而过度变浅。
二、DFlash 的两板斧
DFlash 的全称强调它的两个核心:块扩散草稿(block Diffusion) 和 KV 注入(KV injection)。
1 | 板斧 1:块扩散草稿 → 一次 forward 并行吐出整块 token → 降低 drafting 开销 |
回顾一下投机解码的加速来源,本质就两个因子:
- 接受长度:每个验证周期能接受多少草稿 token,越高越好;
- 草稿相对开销:草稿模型相对 target 多花的成本,越低越好。
绝大多数前作只能改善其中一个因子,DFlash 两个一起改:块扩散降低分母,KV 注入抬高分子。
三、块扩散草稿:一次吐出一整块
自回归草稿的根本问题是「下一个 token 依赖上一个 token」,所以只能串行。块扩散(block diffusion)换了一个生成范式:
1 | 自回归草稿: [t+1] → [t+2] → [t+3] → [t+4] 4 次串行 forward |
它把一个 block 内的若干 token 当作一组「待去噪」的占位,在一次(或极少数几次)forward 里并行地把整块填出来。对 GPU 来说,这是一次高算术强度的稠密计算,正中下怀。
效果有多明显?按 LMSYS 给出的数据,一个 5 层 DFlash 草稿模型,无论生成 4、8 还是 16 个 token,drafting 延迟都比 单层 EAGLE-3 生成 4 个 token 还要低。也就是说,块扩散把「草稿做深」和「草稿做快」这对矛盾解开了——你可以同时拥有更深的草稿模型和更短的 drafting 时间。
直接训练一个小块扩散模型当草稿,接受率会偏低;而直接拿现成的大型扩散 LLM(如 SpecDiff-2)当草稿,显存和 drafting 成本又太高。DFlash 的解法是把块扩散和下面的 KV 注入绑在一起——草稿模型可以很小,因为「理解上下文」这件重活交给了 target。
四、KV 注入:让 target 替草稿理解上下文
Medusa、EAGLE、MTP 的共同直觉是:target 模型本身最懂上下文,所以应该复用 target 的中间表示,而不是让草稿从头再读一遍 prompt。
EAGLE 的做法是把 target 的隐状态喂到草稿模型的输入端。但问题是:这个信号只在第一层注入,越往草稿模型的深层走、越往草稿序列的后面走,target 的上下文信息就衰减得越厉害。
DFlash 把注入的位置换了个地方——直接灌进草稿模型每一层的 KV cache:
1 | EAGLE: target 隐状态 → 草稿模型「输入端」 → 越深越浅、信号衰减 |
具体来说,DFlash 抽取 target 模型对上下文 token 的隐表示,经过一个 KV 投影后,写入草稿模型的 KV cache。这样草稿模型在生成每一块时,注意力都能「回看」由 target 算出来的、高度相关的上下文特征——而且用的就是 target 后面那些层所用的同一批张量。
这带来两个好处:
- 草稿模型可以跳过「从零理解上下文」,专心干一件事——预测下一块 token,因此能做得极小、极快;
- 注入信号不随深度衰减,所以草稿模型可以做得更深、草稿可以拉得更长,而接受率不塌。
一个工程上的小细节:我们不希望把这些 target 隐状态额外存下来占用宝贵的 KV cache 空间,也希望相同前缀的请求能共享 radix cache。所以 DFlash 选择即时物化(immediate materialization)——在草稿前向的其余部分之前,先把草稿的 KV 投影算出来。这一步必须够快,于是 SGLang 为它加了按层批处理的线性投影和一个融合的 Triton kernel(把 norm + RoPE 后处理合并)。
五、为什么这么快:把两个因子拆开看
DFlash 的提速来自块扩散和 KV 注入的叠加,但它们各自的贡献可以通过消融实验分离出来。下面是在同一份数据上、为 Qwen 3-4B 训练的 5 层草稿模型对比,格式为 接受长度 / 端到端加速:
完整 DFlash vs EAGLE-3
| 任务 | EAGLE-3(5 层) | DFlash |
|---|---|---|
| GSM8K | 4.2 / 2.1× | 4.2 / 3.3× |
| HumanEval | 4.3 / 2.2× | 4.0 / 3.2× |
| MT-Bench | 3.1 / 1.4× | 3.0 / 2.2× |
接受长度跟 EAGLE-3 差不多,但靠着超快的并行 drafting,端到端加速明显更高。
只保留块扩散(去掉 KV 注入)
| 任务 | EAGLE-3(5 层) | DFlash(仅块扩散) |
|---|---|---|
| GSM8K | 4.2 / 2.1× | 3.5 / 2.9× |
| HumanEval | 4.3 / 2.2× | 3.5 / 2.9× |
| MT-Bench | 3.1 / 1.4× | 2.6 / 2.0× |
即便接受长度更低,仅靠更快的 drafting,端到端就已经反超 EAGLE-3——这说明降低草稿开销本身就是一个独立的、巨大的杠杆。
只保留 KV 注入(去掉块扩散,回到自回归草稿)
| 任务 | EAGLE-3(5 层) | DFlash(仅 KV 注入) |
|---|---|---|
| GSM8K | 4.2 / 2.1× | 4.8 / 2.4× |
| HumanEval | 4.3 / 2.2× | 4.6 / 2.3× |
| MT-Bench | 3.1 / 1.4× | 3.4 / 1.5× |
接受长度被显著抬高(GSM8K 从 4.2 升到 4.8),端到端也稳超 EAGLE-3——这说明 KV 注入是独立有效的接受率增益。
两个因子各自都能跑赢 EAGLE-3,合起来就是完整 DFlash 的 3.3× / 3.2×。
六、在 SGLang 里落地:从 V1 到 Spec V2
研究阶段的漂亮数字要变成生产可用,需要两步:先在高性能引擎里实现这套机制,再把从 host 调度到 GPU 执行的端到端系统优化干净。DFlash 进 SGLang 也分两步走。
第一步:接入 V1 投机引擎
DFlash 先被加进了(现已废弃的)V1 投机解码引擎,新增了 DFlashWorker(控制草稿模型执行)和它驱动的 DFlashDraftModel。
这里有个反直觉的点值得记住:SGLang 里和 scheduler 对话的是草稿模型 worker(通过 forward_batch_generation 之类的方法),它把 target 模型的 worker 包在里面,草稿就绪后再调用 target 做验证。看代码或 trace 时别搞反了。
DFlash 真正的新东西是 KV 注入把草稿和 target 的状态绑在了一起。对 EAGLE 来说,草稿的 KV cache 是完全私有的、由草稿自己的隐状态投影出来;而 DFlash 是把 target 的隐状态送去做草稿侧的 KV 投影。配合上一节提到的即时物化和融合 kernel,这一步被做得足够快。
第二步:接入 Spec V2 与 overlap 调度
V1 能跑、也快,但还能更快。Spec V2 引擎的核心目标是减少 host-device 同步点——无论 GPU 多快、kernel 多好,频繁的主机-设备同步都会把性能拖垮。解法叫 overlap scheduler(重叠调度器),它抓住两个重叠机会:
1 | 机会 1:批次 N-1 在 GPU 算完后,host 侧的 pop_and_process 清理 |
在 Qwen 3-8B、单卡 B200、并发 32 的设置下,V2 带来了 >33% 的吞吐提升(约 11.4 → 15.3 ktok/s)。这也是为什么 Spec V2 现在成了 SGLang 的默认引擎。
七、实战效果与可用资源
这次随博客发布的是 Qwen3.5-397B-A17B 的 DFlash 草稿模型,在 GSM8K / HumanEval / MT-Bench、并发 1 到 32 的所有设置下,吞吐都高于该模型原生的 MTP 投机:
1 | Qwen3.5-397B-A17B(BF16),HumanEval,贪心解码,思考开启,max new tokens 4096 |
小米新的 MiMo v2.5-Pro-UltraSpeed 也用 DFlash 把单请求输出做到了 >1k tokens/s。
启动一个 DFlash 加速的 SGLang server(节选自原文命令):
1 | export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 |
同样的「块扩散 + KV 注入」思路可以迁移到大多数 target LLM 上,也可以针对自己的数据/模型训练专属的 DFlash 草稿模型。
八、总结
把 DFlash 放进投机解码的演进脉络里看,它的位置很清晰:
| 维度 | 经典 SD | EAGLE / MTP | DFlash |
|---|---|---|---|
| 草稿生成方式 | 小模型自回归 | 轻量草稿自回归 | 块扩散并行 |
| target 上下文复用 | 无 | 输入端注入(会衰减) | 每层 KV 注入(不衰减) |
| 改善的因子 | — | 主要是草稿开销 | 草稿开销 + 接受率 |
| 草稿能否做深 | — | 受开销限制只能浅 | 可深,开销仍低 |
核心结论:
- 投机解码的两个因子可以同时优化——块扩散压低草稿开销,KV 注入抬高接受率,消融实验证明两者各自独立有效;
- 块扩散把「草稿做深」和「草稿做快」解耦,5 层草稿生成 16 token 仍快过单层 EAGLE-3 生成 4 token;
- KV 注入比 EAGLE 的输入端注入更彻底,target 上下文贯穿草稿模型每一层,深层草稿不再失焦;
- 从研究到生产还差一个高性能引擎——Spec V2 的 overlap 调度把 host-device 同步消掉,才让 GPU 的快真正落到端到端吞吐上。
DFlash 给出的启示是:当草稿模型已经被砍到极致之后,下一个杠杆不在「草稿多小」,而在「草稿怎么生成」——把逐 token 的串行自回归,换成对硬件友好的块并行。