DFlash 块扩散解码：以并行扩散替代自回归草稿，实现无损 6 倍推理加速

投机解码的通用框架#

在深入 DFlash 之前，先回顾投机解码的通用流程。

设目标模型为 $\mathcal{T}$，草稿模型为 $\mathcal{D}$，块大小为 $k$。给定前缀序列 $x_{<t}$：

1. 草稿阶段：$\mathcal{D}$ 生成候选序列 $\hat{x}_t, \hat{x}_{t+1}, \dots, \hat{x}_{t+k-1}$。
2. 验证阶段：$\mathcal{T}$ 对 $(x_{<t}, \hat{x}_t, \dots, \hat{x}_{t+k-1})$ 做一次前向传播，得到每个位置的输出分布 $p_{\mathcal{T}}(\cdot \mid x_{<t}, \hat{x}_{<t+i})$。
3. 接受/拒绝：从第一个位置开始，逐一比较草稿 token $\hat{x}_{t+i}$ 与目标模型在该位置的采样结果。找到第一个不匹配的位置 $t+j$，接受前 $j$ 个 token，在该位置重新采样，丢弃后面的所有 token。

关键性质：无论草稿模型如何，验证后的输出分布严格等于目标模型的自回归分布。投机解码是”无损加速”——加速多少取决于草稿模型与目标模型的一致程度，但正确性不会妥协。

实际加速比可以写为：

其中 $t_{\text{target}}$ 是目标模型生成一个 token 的耗时，$t_{\text{draft}}$ 是草稿阶段耗时，$t_{\text{verify}}$ 是验证阶段耗时。$t_{\text{verify}}$ 通常接近 $t_{\text{target}}$（只多了一点注意力计算开销），所以加速比主要由两个因素决定：（1）草稿模型的预测精度（影响有效接受的 token 数），（2）草稿耗时相对于目标模型耗时的比例。

块扩散：从逐 token 到全块并行#

DFlash 的草稿阶段完全抛弃了自回归结构。给定锚点 token（即上一轮目标模型产生的 bonus token），草稿块构建如下：

1. 将锚点 token 放在位置 0。
2. 位置 1 到 $k-1$ 全部填充为 [MASK] token。
3. 整个块通过草稿模型做一次前向传播。
4. 所有 [MASK] 位置的输出 logits 直接用作预测 token。

图源：NVIDIA Technical Blog《Boost Inference Performance up to 15x on NVIDIA Blackwell Using DFlash Speculative Decoding》

草稿模型内部使用双向（非因果）注意力。这意味着位置 3 的预测可以同时看到位置 1、2、4、5 的信息——这在自回归草稿中是不可能的。双向注意力允许模型利用”全局一致性”来生成更连贯的草稿块；例如在生成代码时，如果位置 2 预测了 for，位置 5 可以同时”知道”这个上下文并预测 range，而不需要等位置 2 先生成出来。

但这里有一个微妙之处：双向注意力的”全局一致性”有利有弊。好处是跨位置的信息交换；坏处是各个位置的预测实际上互相影响，可能导致”集体错觉”——一个位置的错误预测通过注意力扩散到其他位置，造成整个块的系统性偏差。这也是为什么后续工作 JetSpec 将注意力改回因果结构的动机之一。

草稿耗时的分析如下。设草稿模型的一次前向传播耗时为 $t_d$，目标模型的一次前向传播耗时为 $t_T$。对于自回归草稿模型（如 EAGLE-3），生成 $k$ 个草稿 token 的耗时为 $k \cdot t_d$；对于 DFlash，无论 $k$ 多大，耗时都是 $t_d$。

取一组典型数据：5 层 DFlash 草稿模型的 $t_d \approx 0.01 t_T$（即目标模型耗时的 1%）。EAGLE-3 单层草稿模型的每步耗时约为 $0.02 t_T$。生成 16 个草稿 token 时，DFlash 只需 $0.01 t_T$，EAGLE-3 需要 $16 \times 0.02 t_T = 0.32 t_T$——差了 32 倍。

这就是为什么即使 DFlash 的单 token 预测准确率低于 EAGLE-3，它的端到端加速比仍然更高：草稿开销低到几乎可以忽略，使得更大的块成为可能。

KV 注入：让每个草稿层直接感知目标模型#

图源：DFlash 论文 Figure 2（arXiv:2602.06036）

这是 DFlash 最关键的设计决策。要理解为什么需要它，先看 EAGLE-3 的做法。

EAGLE-3 将目标模型最后一层的隐藏状态拼接到草稿模型的输入嵌入中。这个设计有一个根本局限：目标模型的特征信号只在输入层注入，经过多个 transformer 层的变换后，信号逐渐被稀释。深层草稿层几乎看不到目标模型的原始特征，只能依赖前面层传来的、经过多次变换的间接信号。

DFlash 的做法完全不同：

1. 从目标模型的多个中间层提取隐藏状态。默认选取 5 层，均匀分布在目标模型的第 2 层到倒数第 3 层之间（例如在 28 层模型中取第 [1, 8, 14, 20, 27] 层）。这些层捕获了不同粒度的语义信息——浅层有更多局部句法特征，深层有更多全局语义特征。

2. 融合提取的隐藏状态。通过一个线性投影层，将来自不同层的隐藏状态（各自维度为 $d_{\text{model}}$）拼接后映射到一个紧凑的”目标上下文特征”向量：

   $$h_{\text{ctx}} = W_f \cdot [h_{l_1}; h_{l_2}; \dots; h_{l_m}] + b_f$$

   其中 $W_f \in \mathbb{R}^{d_{\text{model}} \times (m \cdot d_{\text{model}})}$ 是投影矩阵，$m$ 是选中的目标模型层数。

3. 将融合特征注入每个草稿层的 KV 缓存。对草稿模型的每一层，$h_{\text{ctx}}$ 被分别投影为 Key 投影和 Value 投影，写入该层注意力模块的 KV 缓存中：

   $$K_{\text{ctx}}^{(l)} = h_{\text{ctx}} \cdot W_K^{(l)}, \quad V_{\text{ctx}}^{(l)} = h_{\text{ctx}} \cdot W_V^{(l)}$$

   其中 $l$ 是草稿层索引。在注意力计算中，草稿 token 的 Query 不仅对草稿 token 自身的 Key/Value 做注意力，也对注入的 $K_{\text{ctx}}^{(l)} / V_{\text{ctx}}^{(l)}$ 做注意力：

   $$\text{Attention}(Q_{\text{draft}}, [K_{\text{draft}}; K_{\text{ctx}}], [V_{\text{draft}}; V_{\text{ctx}}])$$

4. 跨轮复用。同一个草稿步骤中，$h_{\text{ctx}}$ 只需要提取一次，然后在所有草稿层、多个草稿轮次中复用（存储在 KV 缓存中）。

这个设计的核心优势：每个草稿层都能直接接触到目标模型的多粒度语义特征。浅草稿层能看到目标浅层的句法信息，深草稿层能看到目标深层的语义信息。这种”深度耦合”远比只在输入层注入一条信号有效。

实验数据也印证了这一点。在 GSM8K 数学推理任务上，同等 5 层草稿模型的条件下：

• 仅输入层融合（EAGLE-3 方式）：平均接受长度 $\tau = 3.5$
• 逐层 KV 注入（DFlash 方式）：平均接受长度 $\tau = 4.2$

20% 的接受长度提升，完全来自特征注入方式的改变。这直接转化为约 20% 的额外加速。

训练策略：随机锚点采样与衰减加权损失#

DFlash 的训练有三个关键设计。

随机锚点采样

推理时，草稿模型需要在任意位置启动扩散过程——本轮可能的锚点 token 可以出现在回复序列的任何位置。训练必须覆盖这种多样性。

具体做法：在每条训练样本的回复序列中，随机采样多个有效位置作为锚点（只采样 loss_mask > 0.5 的位置，排除填充和提示部分）。每个锚点标记一个块的起始位置，其后的 $k-1$ 个位置被 mask 为预测目标。多个块被拼接成一条训练序列，通过 Flex Attention 的稀疏掩码隔离：

• 块内双向可见：同一块内的所有位置可以互相看到（双向注意力）。
• 目标特征始终可见：所有位置都能注意到 KV 注入的 $h_{\text{ctx}}$。
• 块间完全隔离：不同块之间不可见，防止信息泄漏。

每个训练序列最多包含 512 个锚点，这意味着每个样本可以提供多达 512 个独立的草稿训练信号。这种高密度采样极大地提高了训练效率——DFlash 用 EAGLE-3 约 1/5 的训练数据量就超越了它的表现。

图源：DFlash 论文 Figure 4（arXiv:2602.06036）

指数衰减位置权重

投机解码的验证过程是逐位置顺序进行的：一旦在位置 j 发现草稿 token 与目标模型不一致，位置 j 之后的所有 token 都被丢弃。这意味着早期位置的准确率远比晚期位置重要——位置 1 的错误直接浪费整个块，位置 15 的错误只浪费最后一个 token。

DFlash 用一个指数衰减权重来体现这个优先级。对块内第 $i$ 个预测位置（$i = 1, 2, \dots, k-1$，$i=0$ 是锚点不计 loss），权重为：

其中 $\gamma$ 是衰减因子。当 $\gamma$ 较小时（如 $\gamma = 3$），早期位置的权重远大于晚期位置；当 $\gamma$ 较大时（如 $\gamma = 10$），权重分布更均匀。论文中的默认设置倾向于较小的 $\gamma$，以强约束早期位置精度。

最终训练损失为加权交叉熵：

其中 $\hat{p}_i$ 是草稿模型在位置 i 的预测分布，$y_i$ 是目标模型在该位置的真实 token。

D-PACE 变体（Dynamic Position-Aware Cross-Entropy）

标准 DFlash 损失对所有训练样本使用固定的衰减权重。D-PACE 更进一步：动态估计每个位置的学习难度，对困难位置分配更高权重。

具体做法：维护每个位置的平均历史损失 $\bar{\ell}_i$。对当前批次，计算：

如果一个位置的损失显著高于历史平均，说明该位置对当前样本特别困难，应获得更高权重。这使得训练更聚焦于每个样本的”瓶颈 token”——那些草稿模型反复预测错误的关键位置。

验证流程#

验证阶段保持投机解码的标准流程，没有任何变化：

1. 目标模型接收 [前缀; 锚点; 草稿1; 草稿2; ...; 草稿k-1] 的完整序列。
2. 一次前向传播后，得到每个位置的条件概率分布 $p_{\mathcal{T}}(\cdot \mid \text{前缀}, \text{锚点}, \text{草稿}_{<i})$。
3. 从位置 1 开始，逐一比较草稿 token 与目标模型的贪婪采样结果（或温度采样结果）。
4. 找到第一个不匹配的位置，接受之前的所有 token，在该位置用目标模型重新采样作为 bonus token。
5. 返回步骤 1，以新 bonus token 作为锚点开始下一轮。

整个过程保证了最终输出分布严格等于目标模型的自回归分布——这是”无损加速”的含义。

草稿模型架构#

DFlash 草稿模型是一个极轻量级的 Qwen3 风格 decoder-only transformer。一个典型的 5 层草稿模型参数约 340M，主要配置如下：

草稿模型共享目标模型的嵌入层（embed_tokens）和输出投影层（lm_head），且这两部分在训练和推理中都是冻结的。这意味着草稿模型不需要存储两份词表，节省了大量显存。

为什么用 5 层？这是精度与速度的权衡。层数越多，预测精度越高，但草稿耗时也线性增加。论文的消融实验显示：2 层草稿在接受长度上显著弱于 5 层；8 层相比 5 层的接受长度提升很小（<5%），但耗时增加了近 60%。5 层是 experimentally found sweet spot。

目标特征提取与融合#

目标模型中间层特征的提取在目标模型预填充阶段完成，不需要额外的目标模型前向传播。具体流程：

1. 目标模型处理输入前缀，逐层计算直到最后一层。
2. 在预设的中间层位点（如 [1, 8, 14, 20, 27]）处，hook 该层的隐藏状态输出。
3. 将 hook 到的隐藏状态保存到 KV 缓存中，供后续草稿轮次使用。

一个重要的实现细节：选定的是哪些层？ 默认策略是从目标模型的第 2 层到倒数第 3 层之间均匀采样 5 层。为什么跳过了第一层和最后两层？第一层的隐藏状态太接近输入嵌入，语义信息太少；最后两层的隐藏状态已经高度压缩，包含过多任务特定信息，用于草稿模型的通用预测时反而可能引入噪声。

关键伪代码#

以下伪代码展示了 DFlash 推理的核心流程：

1
def dflash_generate(target_model, draft_model, prefix, block_size=16):
2
    # 1. 预填充：目标模型处理前缀
3
    # target_features 是从多个中间层 hook 的隐藏状态
4
    logits, target_features = target_model.prefill(
5
        prefix,
6
        hook_layers=[1, 8, 14, 20, 27]
7
    )
8
    # bonus_token 是预填充的输出 token
9
    bonus_token = sample(logits[-1])
10
    generated = [bonus_token]
11

12
    # 2. 将目标特征融合并注入草稿模型的所有层
13
    context_feature = draft_model.fuse_features(target_features)
14
    draft_model.inject_kv_context(context_feature)  # 注入所有草稿层
15

16
    while not done:
17
        # 3. 构建草稿块：[anchor, MASK, MASK, ..., MASK]
18
        anchor_id = token_to_id(bonus_token)
19
        draft_block = [anchor_id] + [MASK_TOKEN_ID] * (block_size - 1)
20

21
        # 4. 块扩散：一次前向传播预测所有 MASK 位置
22
        draft_logits = draft_model.forward(
23
            draft_block,
24
            attention_mask="bidirectional",  # 块内双向注意力
25
            use_kv_context=True               # 使用注入的目标特征
26
        )
27
        # 取每个 MASK 位置的 argmax 作为草稿 token
28
        draft_tokens = [argmax(logits[i]) for i in range(1, block_size)]
29

30
        # 5. 验证：目标模型一次性检查整个块
31
        full_seq = generated + draft_tokens
32
        target_logits = target_model.forward(full_seq)
33

34
        # 6. 接受/拒绝
35
        accepted = []
36
        for i, draft_tok in enumerate(draft_tokens):
37
            target_tok = argmax(target_logits[len(generated) + i])
38
            if draft_tok == target_tok:
39
                accepted.append(draft_tok)
40
            else:
41
                # 分歧点：用目标模型的 token 作为新 bonus
42
                bonus_token = target_tok
43
                accepted.append(bonus_token)
44
                break
45
        else:
46
            # 全部接受：最后一个位置的预测作为 bonus
47
            bonus_token = sample(target_logits[-1])
48

49
        generated.extend(accepted)
50

51
    return generated

训练流程#

训练数据以对话形式的指令微调数据为主。每条样本包含指令、上下文和期望回复。训练流程：

1. 数据准备：用目标模型对训练数据做一次完整的前向传播，在每个位置的隐藏状态都保存下来（这一步最耗时，但只做一次）。
2. 锚点采样：在每条回复中随机采样多达 512 个锚点位置。
3. 块构建与掩码：为每个锚点构建 [anchor, MASK, ..., MASK] 的草稿块，块大小等于训练块大小（通常 16）。
4. 多块拼接：将所有块按 Flex Attention 稀疏掩码拼接成一条长序列，块间互不可见。
5. 前向+损失：草稿模型前向传播，计算衰减加权交叉熵损失。
6. 优化：AdamW 优化器，学习率 $6\times 10^{-4}$，线性衰减。

训练数据量因模型而异。以 Qwen3-Coder-30B-A3B 的 DFlash 草稿模型为例，训练数据包括：

• NVIDIA Nemotron-Post-Training-Dataset-v2（代码部分）：约 200K 样本
• evol-codealpaca-v1：约 50K 样本
• Cline 执行轨迹：约 39K 样本
• 合计约 289K 样本

相比之下，EAGLE-3 的训练使用了约 1.4M 样本的 open-perfect-blend 数据集。DFlash 用约 1/5 的数据量就超越了 EAGLE-3 的加速比，数据效率显著更高。原因在于：

• 每个样本通过多锚点采样提供了数百个训练信号，而非自回归方法的一个。
• KV 注入提供了更强的条件信号，降低了模型对数据量的依赖。

主要基准测试#

以下数据来自 DFlash 论文和社区验证，涵盖不同模型规模、不同硬件、不同工作负载。

图源：DFlash 论文 Figure 1（arXiv:2602.06036）

Qwen3-8B 上的加速比（greedy decoding）：

代码和数学类任务加速最显著（4-6×），因为这类文本的结构性强、可预测性高，块扩散能充分利用全局一致性。开放域对话的加速较小（2-3×），但仍然是显著的。

Qwen3.6-35B-A3B（MoE 模型，GB10 平台）：

MoE 模型由于稀疏激活的特性，单 token 生成的计算量相对更大，投机解码的相对收益更高。

Gemma 4 26B-A4B（RTX 5090，三种方法的直接对比）：

注意这个表中一个反直觉的数据点：DFlash 的单 token 接受率只有 16%，远低于 EAGLE-3 的 48%，但接受长度（3.45）和端到端加速（2.19×）都更高。这是因为 DFlash 一次并行生成 15 个草稿 token，即使大部分被拒绝，只要平均有 3.45 个被接受，就已经值得了——生成这 15 个草稿的开销几乎和 EAGLE-3 生成 3 个草稿差不多。

按工作负载分解（Gemma 4，DFlash vs EAGLE-3 加速比）：

DFlash 在不同工作负载上的表现差异极大——从散文叙事的几乎无加速，到重复性文本的 4.37 倍加速。这种”冰火两重天”的特征需要在实际部署中认真评估。如果你的应用主要是开放域闲聊，DFlash 可能不是最佳选择；如果是结构化 JSON 输出或代码补全，DFlash 的收益是压倒性的。

NVIDIA Blackwell 上的表现#

DFlash 在 Blackwell GPU 上的优势尤为突出，因为 Blackwell 的设计恰好契合了 DFlash 将串行工作转为并行工作的思想。

在 DGX B300（8× B300 GPU）上使用 TensorRT-LLM + DFlash 部署 gpt-oss-120b 模型：

• 无投机基准：~40 tok/s/用户
• DFlash（k=16）：500-600 tok/s/用户，约 15 倍吞吐量提升
• 相比 EAGLE-3：1.5 倍额外吞吐量提升

这个 15 倍的数字需要理解其条件：它测量的是高并发下的吞吐量增益，不是单用户的延迟改善。当数百个用户同时请求时，DFlash 通过将解码从内存带宽瓶颈转变为计算瓶颈，充分利用了 Blackwell 的 15 PFLOPS NVFP4 算力。单用户 batch size 为 1 时，加速比在 2-4 倍之间。

图源：NVIDIA Technical Blog《Boost Inference Performance up to 15x on NVIDIA Blackwell Using DFlash Speculative Decoding》

在 vLLM 部署中（同等工作负载、相同并发数）：

• Coding 数据集：DFlash 2.6× vs EAGLE-3 1.8×
• RAG 数据集：DFlash 2.3× vs EAGLE-3 1.7×
• Reasoning 数据集：DFlash 2.3× vs EAGLE-3 1.8×

在 SGLang 部署中：DFlash 实现了约 5.1 倍的吞吐量提升。

消融实验#

DFlash 论文中的关键消融实验揭示了几个设计选择的重要性：

KV 注入 vs 输入拼接

取 5 层 DFlash 草稿模型，在 GSM8K 上对比两种特征注入方式：

KV 注入将接受长度提升了 20%。这种优势随草稿模型加深而放大：层数越多，输入层信号被稀释的效应越明显，逐层注入的优势越大。

草稿模型深度

在 Qwen3-8B 上，保持其他配置不变，改变草稿模型层数：

8 层草稿的接受长度几乎不增长（4.2 → 4.3），但草稿耗时近乎翻倍（1.5 → 2.6 ms），导致端到端加速反而下降。这表明草稿模型的深度存在一个”饱和点”，超过之后边际收益为负。

块大小

k=32 时加速开始下降，因为虽然接受长度仍在增长（4.8），但验证阶段需要处理更长序列，验证耗时的增加开始超过额外接受 token 的收益。

推理框架集成#

DFlash 目前支持的推理框架：

vLLM 中的使用方式（来自 vLLM Speculators 文档）：

1
vllm serve Qwen/Qwen3-8B \
2
    --speculative-model z-lab/Qwen3-8B-DFlash \
3
    --speculative-config '{
4
        "method": "dflash",
5
        "block_size": 16,
6
        "target_layer_ids": [1, 8, 14, 20, 27]
7
    }'

SGLang 中的使用方式：

1
python -m sglang.launch_server \
2
    --model Qwen/Qwen3-8B \
3
    --speculative-algorithm dflash \
4
    --speculative-draft-model z-lab/Qwen3-8B-DFlash \
5
    --speculative-block-size 16

支持的模型家族#

截至 2026 年 8 月，Hugging Face 上已有覆盖以下模型族的 DFlash 草稿模型：

• Qwen3.x / Qwen3.5 / Qwen3.6 系列（8B、14B、32B、35B-A3B、Coder 30B-A3B 等）
• Gemma 4 系列（26B-A4B 等）
• Llama 3.1 系列（8B 等）
• gpt-oss 系列（120B 等）
• Kimi K2.6 / K2.7 系列
• Alpamayo-1.5（NVIDIA，自动驾驶 VLA）

训练自己的 DFlash 草稿模型#

NVIDIA 提供了 TensorRT-Model-Optimizer（modelopt>=0.43.0）中的 DFlash 训练工具链：

1
from modelopt.speculative import DFlashTrainer
2

3
trainer = DFlashTrainer(
4
    target_model_path="Qwen/Qwen3-8B",
5
    block_size=16,
6
    num_draft_layers=5,
7
    target_layer_ids=[1, 8, 14, 20, 27],
8
    gamma=3.0,          # 衰减因子
9
    max_anchors=512,    # 每样本最大锚点数
10
)
11
trainer.train(
12
    dataset="your-dataset",
13
    output_path="./my-dflash-drafter",
14
    learning_rate=6e-4,
15
    max_steps=10000,
16
)

训练硬件需求相对亲民：Pilot 品质的草稿模型（5K 样本）可以在单张 RTX PRO 6000 Blackwell（96GB）上完成；生产品质（2M+ 样本）需要多 GPU 训练。

草稿模型的 top-1 匹配率随训练数据量的增长情况：

工作负载敏感性强#

DFlash 在不同类型的文本上表现差异极大。在散文和叙事性文本上，加速比接近 1×（几乎无加速甚至轻微退化），因为这类文本的 token 间依赖关系复杂，块扩散无法有效捕获。在高度结构化的输出（JSON、代码、数学推导）上，加速比可达 4-6×。

这意味着在实际部署中需要做动态策略选择：根据请求类型（或观察到的接受率）在 DFlash、EAGLE-3 和无投机解码之间切换。NVIDIA 的 ModeSwitch-LLM 和 SGLang 的请求级路由已经在探索这个方向。

草稿模型需要逐模型定制#

每个目标模型都需要训练一个对应的 DFlash 草稿模型。草稿模型的隐藏维度、注意力头配置、词表都必须与目标模型对齐。这意味着：

• 不能像 EAGLE-3 那样部分复用草稿架构
• 模型生态中每增加一个新目标模型，都需要额外的草稿模型训练工作
• 微调过的目标模型需要重新训练或至少重新校准草稿模型

高并发场景下的接受率下降#

DFlash 论文中的多数加速数字来自低并发（batch size ≤ 8）或单请求场景。在高并发（batch size ≥ 32）条件下，由于目标模型的 KV 缓存压力增大、批处理效率下降，实际接受的草稿 token 数可能减少。

社区报告显示，在 32 并发时，Qwen3.6-35B-A3B 的 DFlash 加速比从 4.3× 下降到约 2.9×——仍然显著，但下降幅度值得注意。

块扩散的双向注意力问题#

DFlash 的块内双向注意力使得每个位置的预测互相影响。这种”全局一致性”在某些场景下是优势（如 JSON 结构生成），但在需要严格因果关系的场景下会导致系统性错误（如数学证明中的逐步推导）。

这正是后续工作 JetSpec（2026 年 6 月，arXiv:2606.18394）要解决的问题。JetSpec 在 DFlash 的基础上做了两个关键改动：

1. 块内注意力改为因果结构：位置 i 只能看到位置 0 到 i，保证每个位置的预测显式以前面所有位置的预测为条件。
2. 前向 KL 蒸馏损失：用温度缩放的前向 KL 散度替代硬标签交叉熵，让草稿模型学习目标模型的完整软分布，而非仅仅 top-1 token。

JetSpec 将加速比继续推高至最高 9.64 倍。

视觉语言模型支持有限#

当前的 DFlash 实现在视觉语言模型上支持有限。当输入包含视觉信息时，DFlash 的草稿过程被抑制（退化为无投机解码），因为草稿模型无法有效利用视觉特征。这是未来工作的一个方向。