ICML'26 | DFlash：用 Diffusion 模型干掉 AR Drafter，投机解码提速 6×

August 14, 2026

2026   ·   LLM   投机解码   Diffusion   推理优化   论文解读  

ICML’26 | DFlash：用 Diffusion 模型干掉 AR Drafter，投机解码提速 6×

原文：DFlash: Block Diffusion for Flash Speculative Decoding

代码：https://dflash.z-lab.ai/

1. 先说问题：为什么要找一个 drafter

大模型生成文字时，一次只生成一个 token。生成完第 $i$ 个 token，才能把它放回上下文，继续算第 $i+1$ 个 token。这种“前一个 token 决定后一个 token”的生成方式叫 autoregressive（AR，自回归）。

AR decode 的问题不是显卡不会算，而是每一步都要重新经过模型。生成 100 个 token，就要连续做 100 次依赖前一步结果的 forward，GPU 很难把这些计算合并起来。

投机解码（speculative decoding） 的思路是把工作拆给两个模型：一个小模型先猜未来几个 token，这个小模型叫 draft model / drafter；真正的大模型再把这段猜测一次性拿来验证。猜对的 token 直接接受，猜错的位置从那里截断，最后仍然由大模型决定输出。因此，只要验证规则正确，最终分布和大模型自己逐 token 生成完全一致。

一轮投机解码的平均每 token 延迟可以粗略写成：

\[L = \frac{T_\text{draft} + T_\text{verify}}{\tau}\]

$T_\text{draft}$ 是小模型打草稿的时间，$T_\text{verify}$ 是大模型批量验证的时间，$\tau$ 是每轮平均接受的 token 数。过去的工作主要想办法提高草稿质量，让 $\tau$ 变大；DFlash 抓住的是另一个瓶颈：AR drafter 自己也在一个 token 一个 token 地生成，$T_\text{draft}$ 随草稿长度线性增长。

2. EAGLE-3：以前的 drafter 是怎么工作的

理解 DFlash 前，先看它要对标的方法，EAGLE-3： 它是一种基于目标模型 hidden feature 的 AR drafter。这里的 hidden state / hidden feature 可以先理解成目标模型在某一层对当前上下文形成的一串内部向量。它不是“下一个 token 的概率”，而是模型读完上下文后留下的中间表示，里面包含语义、句法以及对后续 token 的预测信息。EAGLE-3 把目标模型的 hidden feature 交给一个更小的 drafter，帮助小模型不要从零猜。

EAGLE-3 的推理过程仍然是 AR 的。假设目标模型已经确认了：

The capital of France is

EAGLE-3 会这样草拟：

第 1 次 drafter forward：预测 Paris
第 2 次 drafter forward：把 Paris 接回输入，再预测 .
第 3 次 drafter forward：再预测 <eos>

为了提高接受率，EAGLE-3 通常不只生成一条序列，而是生成一个 draft tree（草稿树）：每个节点代表一个候选 token，多个候选路径一起交给目标模型验证。例如 EAGLE-3 (60) 表示草稿树有 60 个节点，不是说模型有 60 层。目标模型可以用 tree attention 一次检查这些候选路径。

树让候选更多，但没有改变最根本的成本：树里的候选仍然需要通过多次串行的 drafter forward 生成。想草拟得更长，就要付出更多次 forward；为了控制延迟，EAGLE-3 的 drafter 通常只能做得很浅。这就是 DFlash 想解决的具体问题。

3. DFlash 先看完整的一轮流程

DFlash 不是一个新的基础语言模型，而是一套把 diffusion drafter 接入投机解码的框架。它没有把 Qwen3 改造成一个新的 diffusion 主模型：原来的目标 AR 模型 $M_t$ 保持不变，另外训练一个很小的 diffusion drafter $M_d$。目标模型负责最终答案和验证，drafter 只负责猜一小段。

为了描述这段流程，先定义两个词：

  • 已确认前缀：到当前时刻为止，已经被目标模型接受的 token 序列。
  • anchor：已确认前缀的最后一个 token。下一轮草稿从它后面开始。

DFlash 把“anchor 后面要猜的固定长度窗口”叫一个 block。论文常用 block_size=16：第 1 个位置放 anchor，后面 15 个位置留给 drafter 预测。

一轮完整的 DFlash 投机解码如下：

已确认前缀：[x₁, x₂, ..., xₙ]
                         ↑
                       anchor xₙ

1. 目标模型在第一次完整处理 prompt、以及后续验证新 token 时产生 hidden states，抽出并缓存 context feature
2. drafter 输入：xₙ + <mask> <mask> ... <mask> + context feature
3. drafter 一次 forward，同时输出后面 15 个位置的 token 分布
4. 从这 15 组分布得到 draft tokens
5. 目标模型把 draft tokens 接到前缀后面，一次 forward 批量验证
6. 接受从前往后连续正确的 token，并让目标模型额外产生一个 bonus token（验证结束后顺手得到的下一个 token）
7. 用新的已确认前缀和 bonus token 开始下一轮

这里“下一轮”是沿着整段回答向前移动，不是对同一个 block 再做第二次 denoising。DFlash 的 block 内 denoising 步数就是 1。

目标模型在第 1 步和第 5 步都在工作，但两次工作的目的不同：第 1 步提供上下文表示，第 5 步负责最终验收。context feature 不是一个凭空出现的变量，它就是目标模型选定若干层的 hidden state 拼接、投影后的结果。

4. Diffusion 到底体现在哪里

4.1 不是 AR 的“下一个 token”，而是被 mask 的一段 token

普通 AR drafter 的输入是一个已经完整的前缀，输出下一个 token；DFlash 的输入在待预测区域放的是 <mask>：

真实序列：       The capital of France is Paris . <eos>

送进 drafter：   The capital of France is <mask> <mask> <mask>
训练标签：                                      Paris     .       <eos>

模型不是只在最后一个位置计算“下一个 token 是什么”，而是对三个 mask 位置同时计算三组词表 logits。logits 可以理解成 softmax 之前的词表分数，分数最高的 token 就是最可能的候选。这种“把一段真实序列破坏成 mask，再恢复原 token”的建模方式，就是 DFlash 借用的 masked diffusion / block diffusion 思路。

4.2 它有没有 diffusion 的迭代步数

常见的 diffusion LLM 往往会反复 denoise 同一个序列：

第 0 步：<mask> <mask> <mask> <mask>
第 1 步：回填一部分 token，剩余位置仍是 <mask>
第 2 步：基于已回填 token，继续预测剩余 mask
...
第 S 步：全部位置变成真实 token

每一步都需要一次模型 forward。多做几步，后面的位置可以看到更多已经确定的 token，通常更准，但延迟也更高。

DFlash 的推理配置相当于把同一个 block 的 denoising 步数设为 1。它不回填后再次 forward，而是一次性预测整个 mask block，然后立即交给目标模型验证。回答变长时，系统会不断开启新的 block；这是一轮一轮的 speculative decoding，不是同一 block 内的多步 diffusion。

因此，从代码行为看，DFlash 的 diffusion drafter 可以直接概括为：

mask block --(一次联合 forward)--> 每个位置一组 logits
                                      │
                                      └--> 每个位置采样一个 token

不是分别跑三次模型。一次 forward 联合地产生三组 logits，三组 logits 都读取同一个 target context 和同一个 mask block；forward 结束后，才分别采样出三个 draft token。然后目标模型验证整段草稿。

这个方案牺牲了什么？预测第三个位置时，模型还没有看到自己采样出的前两个 token，所以 block 内没有 AR 那种逐 token 依赖。它换来的收益是：草拟 15 个 token 只需要一次 drafter forward，而不是 15 次串行 forward。

5. Diffusion drafter 的结构：代码比名词更直观

DFlash 的 drafter 不是一个 7B 的完整 diffusion LLM，而是一个很小的 Transformer。论文主配置使用 5 层 drafter，Qwen3-Coder 使用 8 层；token embedding（把 token ID 变成向量的查表层）和 LM head（把最后的 hidden state 映射回词表分数的输出层）复用目标模型并冻结，真正训练的是 drafter Transformer 层和 context projection。

目标模型先把若干层 hidden state 拼起来，压到 drafter 的 hidden size：

\[H_t=\text{RMSNorm}\left(W_c[H^{(l_1)};\ldots;H^{(l_5)}]\right)\]

$H_t$ 就是上面流程中的 target context。它对应已确认前缀的每个位置，并在后续 verification 时追加新确认 token 的 feature。

下面是一个删掉工程细节后的 forward。形状比具体 API 更重要：

class DFlashLayer(nn.Module):
    def forward(self, h_block, h_ctx):
        # h_block: [B, 1 + gamma, D]
        #   第 0 个位置是 anchor，后 gamma 个位置是 mask
        # h_ctx:   [B, N, D]
        #   目标模型已确认前缀的 context feature

        # block token 产生 Query
        q = self.wq(h_block)                       # [B, 1+gamma, D]

        # context 和 block token 一起作为 Key / Value
        k = torch.cat([self.wk(h_ctx),
                       self.wk(h_block)], dim=1)   # [B, N+1+gamma, D]
        v = torch.cat([self.wv(h_ctx),
                       self.wv(h_block)], dim=1)   # [B, N+1+gamma, D]

        # block 内双向可见；不能看到未来 block
        h = attention(q, k, v, mask="block-bidirectional")
        return self.ffn(self.wo(h))


def draft_one_block(anchor_id, context, drafter, gamma):
    # 输入：1 个已确认 anchor + gamma 个 mask
    ids = torch.cat([
        anchor_id[:, None],
        torch.full((batch, gamma), MASK_ID)
    ], dim=1)                                      # [B, 1+gamma]

    h = drafter.embedding(ids)
    for layer in drafter.layers:
        h = layer(h, context)

    # 丢掉 anchor 位置，只读取 mask 位置的 logits
    logits = drafter.lm_head(h[:, 1:, :])          # [B, gamma, vocab]
    draft_tokens = sample(logits)                  # [B, gamma]
    return draft_tokens

这段代码里 diffusion 的位置非常具体：输入有一串 <mask>，同一个 block 内的 mask token 通过双向 attention 一起处理，最后一次性输出所有 mask 位置的 logits。sample(logits) 只负责从每个位置的分布取 token；它不会再调用 drafter，所以这里没有隐藏的第二轮 denoising。

target context 的注入是 DFlash 的第二个关键点。普通做法可以把目标模型 feature 拼到 drafter 的输入上，但层数变深后信息容易被稀释。DFlash 把 context 直接放进每一层 attention 的 Key、Value，因此每层都能访问目标模型的内部表示：

\[Q_i=W_i^QH_d,\quad K_i=[W_i^KH_t;W_i^KH_d],\quad V_i=[W_i^VH_t;W_i^VH_d]\]

DFlash 推理设计：context feature 注入 KV cache

6. 训练：它不是拿来即用的通用 diffusion 模型

DFlash 的 drafter 需要针对目标模型训练。目标模型冻结，drafter 学的是：“在这个目标模型给出的 hidden context 下，如何一次预测下一个 block”。换一个目标模型，通常要重新准备数据、重新抽 hidden feature，并训练对应的 drafter。

6.1 训练样本怎么构造

论文主实验使用约 80 万条样本，混合 Nemotron Post-Training Dataset V2 和 CodeAlpaca。作者先让目标模型自己根据 prompt 生成 response，再用这些 target-generated response 训练 drafter。这样 drafter 对齐的是目标模型的输出分布，而不是数据集原始答案的分布。

一条样本的构造过程是：

目标模型生成的 clean sequence：prompt [r₁ r₂ r₃ r₄ r₅ r₆ ...]

随机选 anchor：                prompt [r₁ r₂] [r₃ r₄ r₅ r₆]
                                                ↑
                                              block

drafter 输入：                 prompt [r₁ r₂] [<mask> <mask> <mask> <mask>]
训练标签：                                      [r₃     r₄     r₅     r₆]

第一个位置 r₂ 是干净的 anchor，后面四个 token 被 mask，drafter 同时预测四个位置。论文的 block_size 包含 anchor，所以 block_size=16 时，实际预测的是 15 个 token。

目标模型先跑完整的 clean sequence，抽取选定层的 hidden states；这些 hidden states 经过 $W_c$ 后作为 drafter attention 的 context KV。训练时目标模型不反向更新。

6.2 更新哪些参数，训练规模多大

论文实现复用目标模型的 token embedding 和 LM head，并保持它们冻结；更新的是 5 层左右的 drafter Transformer 和 context projection $W_c$。主配置是 6 个 epoch、AdamW、学习率 $6\times10^{-4}$、gradient clipping 1.0、cosine schedule、warmup ratio 0.04；最长序列 3072 token，Qwen3-Coder 使用 4096 token，每条序列每个 epoch 随机采样 512 个 anchor。

80 万是样本数，不是一个普适的最低 token 数。论文还用 100K 样本做消融，给长上下文模型用 1.6K 条 LongAlign-10K 样本 fine-tune 3 个 epoch。真正的成本不只是样本数量：每个目标模型都要生成对齐 response，并用目标模型跑一遍抽 hidden feature。

target feature 有两种处理方式：

  • online：训练每一步现跑目标模型，省磁盘但增加计算；
  • offline：提前跑目标模型，把 hidden feature 缓存下来，训练 drafter 时直接读取，省计算但占存储。

如果抽 5 层、序列很长、样本很多，offline cache 会迅速膨胀；如果用 online，又必须承担目标模型 forward 的训练开销。这是复现时比“写一个 mask loss”更现实的门槛。

6.3 Loss 和 attention mask

每个 mask 位置都计算 token-level cross-entropy，但越靠前的位置权重越高：

\(\mathcal{L}=-\sum_{k=1}^{\gamma}w_k\log p_d(x_k^*)\) \(w_k=e^{-(k-1)/\gamma}\)

原因是验证从前往后进行：第一个草稿 token 错了，后面即使猜对，也不能形成连续接受前缀。

训练时可以把多个 block 拼成一条长序列，但 attention mask 必须限制信息流：同一个 block 内 token 双向可见，不同 block 之间不能互相看到未来。否则训练时模型会偷看推理阶段不存在的答案。

这套训练有三个实际难点：

  1. 目标模型绑定。数据、tokenizer、hidden representation 和 drafter 是配套的，不能把一个 Qwen3 drafter 随意接到别的架构上。
  2. 长上下文成本。hidden feature 的缓存量随序列长度和抽取层数线性增加；online 方案则把目标模型计算成本放回训练。
  3. block size 不是免费可调的。训练时见过的 block 长度会影响 acceptance。论文观察到大 block 训练的模型缩小 block 通常还能工作，反过来用小 block 训练的模型推理时突然放大 block 则不稳定。

所以 DFlash 不是“加一个 mask token 就能复现”的技巧，而是一套目标模型、数据生成、hidden feature 缓存、drafter 训练和推理运行时共同配合的系统。

7. 为什么它能比 AR drafter 快

AR drafter 草拟 16 个 token，需要多次串行 forward；DFlash 的 16-position block 只需一次 drafter forward。这个差别让 DFlash 可以在相同 latency budget 下使用更深的 drafter：论文实测，5 层 DFlash 草拟 16 个 token 约 5ms，而 1 层 EAGLE-3 草拟 16 个 token 约 26ms。

EAGLE-3 vs DFlash drafting 延迟对比

DFlash 并不是因为单个位置预测得比 AR 更准才赢，而是把节省下来的串行时间换成了更大的 drafter 和更长的 block。AR drafter 的问题是质量和延迟绑死；DFlash 用 block 内并行把这两个量拆开了。

在 Qwen3-4B/8B、thinking mode 关闭的实验中，DFlash 平均达到约 4.9× greedy speedup，部分 Math 任务 acceptance length $\tau$ 达到 7.87；论文报告的最高端到端加速超过 6×，并且超过 EAGLE-3 的大树配置。

8. 最后把边界说清楚

DFlash 的核心不是“diffusion 神奇地一次生成任意长度文本”，而是一个非常具体的工程折中：

  • 用 mask block 替代 AR drafter 的逐 token 草拟；
  • 用一次联合 forward 预测 block 内所有位置；
  • 用目标模型 hidden feature 弥补 block 内没有 AR 依赖的问题；
  • 用目标模型 verification 保证最终输出仍然正确。

因此它更准确的名字应该理解成：一个 one-step block-diffusion drafter，用于 speculative decoding。它保留了 diffusion 的 mask-and-denoise 形式，但没有采用完整 diffusion LLM 常见的多轮 denoising 生成。速度来自这个取舍，训练和复现的难点也来自这个取舍。

如果这篇文章涉及的 LLM 推理效率优化你想系统深入，可以看看我之前出版的《动手学 AutoML：从 NAS 到大模型优化实战》，书里有专章讲 LLM 推理效率和参数高效微调，和本文的工程背景有直接关联。

动手学AutoML书籍封面

Flag Counter