DFlash: Block Diffusion for Flash Speculative Decoding

作者/机构: Jian Chen, Yesheng Liang, Zhijian Liu

A1 主要贡献

核心问题与研究目标: 尽管大型语言模型（LLMs）性能强大，但其自回归（Autoregressive）生成方式是逐个token进行的，这导致了固有的序列化解码瓶颈，使得推理延迟高、GPU利用率低。随着链式思考（Chain-of-Thought, CoT）等长推理模型的出现，这一问题愈发严重。现有主流解决方案是推测解码（Speculative Decoding），它使用一个轻量级的草稿模型来推测未来的token序列，然后由目标LLM并行验证。然而，当前最先进的方法如EAGLE-3仍依赖于自回归的草稿生成方式，这种方式效率低下且容易累积错误，将实际加速比限制在约2-3倍。另一方面，扩散式LLMs（Diffusion LLMs）支持并行生成，但其生成质量通常不如自回归模型，且为了保证质量需要较多的去噪步骤，从而降低了原始推理速度。因此，本文的核心研究问题是：能否结合自回归模型的高质量和扩散模型的并行生成能力，同时规避它们各自的缺点？

创新点与核心贡献:
本文提出了DFlash，一个采用轻量级块扩散模型（block diffusion model）进行并行草稿生成的推测解码框架。其核心洞见是“目标模型最了解情况”（the target knows best），即大型自回归LLM的隐藏层特征隐式地包含了关于未来多个token的信息。DFlash利用这一洞见，将草稿模型构建为一个扩散适配器（diffusion adapter），通过以下方式实现高效且高质量的草稿生成：

  1. 基于目标模型上下文特征的条件化生成: DFlash从目标模型的隐藏层中提取深层上下文特征，并将这些特征作为条件注入到草稿模型中。这使得轻量级的草稿模型不必从零开始推理，而是能够有效利用目标模型的强大建模能力来并行预测未来的token块。
  2. 并行扩散草稿生成: 与自回归草稿模型逐个生成token不同，DFlash的草稿模型在一次前向传播中并行生成一整个token块。这种方式极大地降低了草稿生成的延迟，并提高了GPU的利用率。
  3. 高效的训练与推理设计: 通过KV注入（KV injection）、随机掩码块采样（random sampling of masked blocks）和损失加权（loss weighting）等训练策略，DFlash能够高效地将草稿模型与目标模型对齐。

主要成果: 实验证明，DFlash在一系列模型和任务上实现了超过6倍的无损加速，比当前最先进的推测解码方法EAGLE-3快了高达2.5倍。这表明DFlash在加速LLM推理方面取得了显著进步。

图1. DFlash、EAGLE-3与自回归解码在Qwen3-8B上的加速比比较。总体而言，DFlash的加速比是EAGLE-3的2.5倍以上。
图1. DFlash、EAGLE-3与自回归解码在Qwen3-8B上的加速比比较。总体而言，DFlash的加速比是EAGLE-3的2.5倍以上。

A3 背景知识

本节规范了推测解码的加速机制，并阐明了自回归与基于扩散的草稿生成之间的效率权衡。我们的分析强调了为什么扩散草稿模型能够独特地同时实现低草稿延迟和高接受率。

3.1 推测解码的加速

推测解码加速原理。推测解码利用一个较小的草稿模型$M_d$来加速目标模型$M_t$的推理。在每个解码周期中，草稿模型提出$\gamma$个token，然后由目标模型并行验证。遵循Sadhukhan等人【26，MagicDec: Breaking the latency-throughput tradeoff for long context generation with speculative decoding，2025】的分析，平均每个token的延迟为：

$$L = \frac{T_{\text{draft}} + T_{\text{verify}}}{\tau},$$

其中，$T_{draft}$是生成草稿token的时间，$T_{verify}$是验证的成本，而$\tau \in [1, \gamma + 1]$是每个周期预期的接受token数量，包括目标模型产生的额外token。设$L_{target}$为$M_t$的自回归单token延迟；最终的加速比为$\eta = L_{target}/L$。这个表达式明确了其中的权衡：要么通过增加预期的接受长度$\tau$，要么通过减少草稿开销$T_{draft}$来提高加速比。

3.2 自回归与扩散式草稿生成

自回归草稿生成的成本。自回归草稿模型顺序生成token，产生的草稿成本为：

$$T_{\text {draft }}=\gamma \cdot t_{\text {step }},$$

其中$t_{step}$是单次前向传播的延迟。因此，草稿成本随着推测预算$\gamma$线性增长。为了将延迟控制在可接受的范围内，自回归草稿模型被限制在非常浅的架构上（例如EAGLE-3中的单层transformer）。这严重限制了草稿质量：虽然增加$\gamma$会增加草稿成本，但由于模型容量有限，接受长度$\tau$很快就饱和了。在实践中，这种不平衡限制了可达到的加速效果。

扩散草稿生成的成本。扩散草稿模型在一次前向传播中并行生成所有$\gamma$个token，得到：

$$T_{\text{draft}} = t_{\text{parallel}},$$

其中$t_{parallel}$表示块生成的延迟。现代GPU执行此类并行操作的效率远高于多次顺序传播，使得对于同等规模的模型，$t_{parallel} \ll \gamma \cdot t_{step}$。对于中等大小的块，因此$T_{draft}$对$\gamma$基本不敏感。

图3. 1、3、5层DFlash和1层EAGLE-3的草稿成本。
图3. 1、3、5层DFlash和1层EAGLE-3的草稿成本。

设计空间的根本改变。这种并行性从根本上改变了设计空间。由于草稿成本不再随着生成token的数量扩展，扩散草稿模型可以采用更深、更具表达力的架构而不会牺牲延迟。这种增加的容量显著提高了草稿质量和接受长度。根据经验，一个生成16个token的五层DFlash草稿模型，相比于生成8个token的EAGLE-3，既实现了更低的延迟（如图3所示），也获得了更高的接受长度，这使得DFlash在草稿质量和草稿成本之间处于一个更有利的帕累托前沿。

A2 方法细节

4.1 推理

DFlash系统设计。DFlash的系统设计如图2所示。本节将解释DFlash如何利用一个非常小而高效的草稿模型实现高草稿接受长度的关键设计选择。

图2. DFlash推理设计。从目标模型中提取的隐藏上下文特征被融合并注入到每个草稿层的键值缓存（Key-Value cache）中，以实现条件化推测。
图2. DFlash推理设计。从目标模型中提取的隐藏上下文特征被融合并注入到每个草稿层的键值缓存（Key-Value cache）中，以实现条件化推测。

来自目标模型的上下文特征。先前的工作，如An等人【1，PARD: Accelerating llm inference with low-cost parallel draft model adaptation，2025】的工作，简单地将一个小型扩散模型用作推测草稿模型，这导致了较差的接受长度和有限的加速。为了验证这一点，我们训练了一个五层的块扩散草稿模型，它没有任何来自目标模型的条件化，并在几个数学基准上进行了评估。如表8所示，由此产生的加速效果很一般，通常在2-3倍左右。这一限制源于缺乏丰富的上下文指导：在无法访问目标模型内部表示的情况下，扩散草稿模型必须有效地从头预测未来的token。相比之下，大型自回归目标模型的隐藏表示编码的信息远多于token级别的logits。这些特征捕捉了长程依赖和任务特定的语义，并且——至关重要的是——隐式地编码了关于未来token预测的信息，这一点也被Samragh等人【27，Your llm knows the future: Uncovering its multi-token prediction potential，2025】观察到。

特征提取与融合。在DFlash中，给定一个输入提示，目标模型首先执行一次标准的预填充（prefill）遍，以生成第一个token。在此过程中，我们从一组从浅到深均匀采样的固定层中提取隐藏表示。这些隐藏状态被连接起来，并通过一个轻量级的投影层，将跨层信息融合成一个紧凑的目标上下文特征，然后用于条件化草-稿模型。

通过KV注入实现接受率扩展。现有的方法如EAGLE-3也利用目标模型的隐藏特征，但它们将这些特征与草稿模型的token嵌入融合，并仅将其作为输入提供给草稿模型。随着草稿模型深度的增加，来自目标模型的信息变得越来越稀释，导致增加更多草稿层时接受长度的增益递减。DFlash采取了根本不同的策略。我们将融合的目标上下文特征视为持久的上下文信息，并将其直接注入到每个草稿模型层的Key和Value投影中。投影后的特征存储在草稿模型的KV缓存中，并在草稿迭代中重复使用。这种设计为整个草稿模型提供了强大而一致的条件化，使得接受长度能够随着草稿层数的增加而有效扩展。我们将在第5.4.2节中更详细地分析这种行为。

并行扩散草稿生成。DFlash速度的另一个关键贡献者是其低草稿延迟。自回归草稿模型必须执行多次顺序的前向传播来生成草稿token或树，这限制了并行性并导致GPU利用率低下。相比之下，DFlash使用块级扩散过程来预测下一个token块。块内的所有被掩码的位置都在一次前向传播中并行解码。与自回归草稿生成相比，这种块级并行生成显著降低了草稿延迟，并实现了更高的硬件利用率，即使在使用更深的草稿模型时也是如此。

DFlash核心优势总结。总的来说，DFlash将基于扩散的并行草稿生成与来自目标模型的紧密耦合条件化相结合，从而以显著降低的草稿延迟实现了高质量的草稿生成。

4.2 训练

训练目标。DFlash草稿模型被训练用于使块级扩散预测与一个固定的自回归目标模型的输出对齐。我们没有直接采用标准的块扩散训练方法（如Arriola等人【2，Block diffusion: Interpolating between autoregressive and diffusion language models，2025】），而是引入了几个关键的修改，以提高训练效率、可扩展性，并更好地与推理时的推测解码行为对齐。

训练中的KV注入。遵循推理流程，给定一个由提示和其响应组成的序列，我们首先将整个干净的序列通过目标模型，以提取并融合所有token的隐藏特征。然后，如图4所示，这些隐藏特征作为Key和Value投影被注入到草稿模型中。

图4. DFlash训练注意力机制。目标模型提供上下文特征（蓝色）来条件化草稿模型。输入由干净的提示token p和干净的响应token r组成。在每个掩码块内，从干净的响应token中随机抽取一个子集作为锚点（黄色），而掩码token m（绿色）标记了用于并行预测的位置。不可见token（白色）表示注意力掩码，它强制执行因果一致性并防止训练期间块间信息泄露。
图4. DFlash训练注意力机制。目标模型提供上下文特征（蓝色）来条件化草稿模型。输入由干净的提示token p和干净的响应token r组成。在每个掩码块内，从干净的响应token中随机抽取一个子集作为锚点（黄色），而掩码token m（绿色）标记了用于并行预测的位置。不可见token（白色）表示注意力掩码，它强制执行因果一致性并防止训练期间块间信息泄露。

掩码块的随机采样。在标准的块扩散训练中，响应被均匀地划分为块，并对每个块内的随机位置进行掩码，模型被训练来去噪这些被掩码的token。DFlash则针对推测解码的场景定制了块的构建方式。我们从响应中随机采样锚点token，将每个锚点用作一个块的第一个位置，并掩码其余的位置。草稿模型被训练来并行预测接下来的block size - 1个token。这直接匹配了推理时的行为，即草稿模型总是以目标模型产生的一个干净token为条件（即上一步验证产生的额外token）。随机化锚点位置也使草稿模型接触到更多样化的目标上下文特征，提高了数据效率和覆盖范围。如表9所示，这种策略显著提高了接受长度和加速比。

多块并行训练。在训练期间，所有块被连接成一个单一序列，并使用如图4所示的稀疏注意力掩码进行联合处理。在同一块内，token之间以及与注入的目标上下文特征之间进行双向注意力计算，而不同块之间的注意力则被禁止。这种设计使得可以使用Flex Attention【10，Flex attention: A programming model for generating optimized attention kernels，2024】在单次前向和后向传播中高效地训练多个草稿块。

高效的长上下文训练。对于像EAGLE-3这样的方法来说，在长上下文上训练推测草稿模型是具有挑战性的，因为它们在训练时进行测试的成本很高。DFlash通过固定每个序列的掩码块数量，并在每个epoch为每个序列随机采样锚点位置，实现了高效的长上下文训练。这种策略在保持训练成本有限的同时，提供了有效的数据增强。

加速收敛的损失加权。在推测解码中，并非所有token都同等重要。草稿块中早期位置的错误会使所有后续的token都失效。这使得早期的预测对接受长度尤为重要。我们通过对交叉熵损失进行加权来反映这种不对称性，从而在训练期间强调早期的token位置。具体来说，对于块中位置为k的token，我们应用一个指数衰减的权重：

$$w_k = \exp\left( - \frac{k - 1}{\gamma} \right),$$

其中$\gamma$控制衰减率。这种加权优先考虑早期位置，加速了收敛，并比均匀加权产生了更高的接受长度（图5）。

共享嵌入层和语言模型头。为了提高训练效率，草稿模型与目标模型共享token嵌入层和语言模型头，并在训练期间保持它们冻结。只有草稿的Transformer层被更新。这种设计减少了可训练参数的数量，并促使草稿模型作为一个与目标模型表示空间紧密对齐的轻量级扩散适配器。

A4 实验环境

  • 模型:

    • LLaMA-3.1 Instruct (8B)
    • Qwen3 (4B, 8B, Coder-30BA3B-Instruct)

  • 数据集与任务:

    • 训练数据: 混合了约80万个来自NVIDIA Nemotron Post-Training Dataset V2 【22】 和 CodeAlpaca 【5】 的样本。为了更好地与目标模型对齐，训练集中的响应由目标模型重新生成。
    • 评估任务:
      • 数学 (Math): GSM8K 【8】, MATH 【19】, AIME25 【21】
      • 代码 (Code): HumanEval 【6】, MBPP 【3】, LiveCodeBench 【13】
      • 对话 (Chat): MT-Bench 【35】, Alpaca 【30】

  • 硬件配置:

    • 主要实验在NVIDIA H200 GPU上进行。
    • SGLang性能评估在单张NVIDIA B200 GPU上进行。

  • 软件配置:

    • 代码实现: DFlash草稿模型层数为5层（Qwen3 Coder为8层），块大小为16（LLaMA 3.1为10）。目标模型的隐藏特征从第2层到倒数第3层之间均匀选取的5个层中提取。
    • 依赖库与框架: 使用Transformers后端和SGLang 【36】 (搭配FlashAttention-4或Flashinfer后端) 进行评估。
    • 基线模型: Vanilla自回归解码和EAGLE-3 【17】。

A4 实验结果

5.1 指令微调模型（Instruct Models）

在禁用思考模式（thinking mode）的Qwen3模型上，使用Transformers后端对DFlash和EAGLE-3进行了比较。如表1所示，DFlash在所有任务和设置中均显著优于EAGLE-3。

  • 贪婪解码 (Temperature=0): DFlash的平均加速比达到4.9倍，比EAGLE-3（草稿树大小为16）高出2.4倍。
  • 非贪婪采样 (Temperature=1): DFlash的平均加速比为4.1倍，比EAGLE-3高出2.2倍。
  • 与大草稿树比较: 即便与草稿树大小为60的EAGLE-3相比，DFlash在实现更高接受长度的同时，验证开销也显著更低。

这些结果证明了DFlash中基于扩散的草稿生成方式的高效性。

表1. Qwen3模型在禁用思考模式和最大生成2048个token的情况下，相对于基线的解码加速比和平均接受长度(τ)。括号中的数值分别表示EAGLE-3的草稿树大小和DFlash的扩散块大小。
图片描述

5.2 推理模型（Reasoning Models）

在启用思考模式的Qwen3模型上评估DFlash，草稿模型在包含推理轨迹的目标模型输出上进行训练。如表2所示，DFlash依然保持了高接受长度，相对于基线分别实现了约4.5倍和3.9倍的加速比。这一效率提升对于推理时间较长的推理模型的实际部署尤为重要。

表2. 启用思考模式下的解码加速比和平均接受长度(τ)。
图片描述

5.3 在SGLang上的性能

在流行的开源推理框架SGLang上，使用B200 GPU和FA4后端对DFlash进行了实际服务场景的性能评估。如表3所示，在Qwen3-4B、8B和Coder-30B-A3B模型上，DFlash在1到32的并发水平下均提供了稳定的加速，在Qwen3-8B上最高达到了5.1倍的加速。这证明了DFlash在真实世界服务场景中能够显著降低服务成本的实用价值。

表3. 在SGLang (FA4后端)上的吞吐量(tok/s)、相对于基线的加速比和平均接受长度τ。
图片描述

5.4 消融研究

5.4.1 训练数据的影响

在LLaMA-3.1-8B-Instruct模型上，使用与EAGLE-3完全相同的训练数据（UltraChat和ShareGPT）训练DFlash。如表4所示，在SGLang上，DFlash在所有任务、并发水平和EAGLE-3配置下都持续优于EAGLE-3，这证明了DFlash相对于自回归树状推测解码的鲁棒性和效率优势。

表4. LLaMA-3.1-8B-Instruct在SGLang (Flashinfer后端，单B200 GPU)上的加速比和平均接受长度τ。基线报告了绝对吞吐量(TPS; tokens/s)。
图片描述

5.4.2 草稿模型层数

实验探究了草稿模型层数对性能的影响。如表5所示，虽然8层模型能实现更长的接受长度（更高的草稿质量），但5层模型因在草稿成本和质量之间取得了更好的平衡，从而获得了更高的整体加速比。这表明最优层数取决于具体的部署设置。

表5. 5层草稿模型具有最佳的平均加速比。
图片描述

5.4.3 目标模型隐藏特征数量

从更多目标模型层数中提取特征能提供更丰富的语义信息，从而提高草稿质量。如表6所示，使用5个隐藏特征的性能始终优于使用3个。然而，这会带来更高的训练成本，因为缓存目标隐藏状态所需的存储空间会线性增加。

表6. 更多来自目标模型的隐藏特征可以增加接受长度。
图片描述

5.4.4 训练与推理时的块大小

研究发现，在训练和推理时使用更大的块大小（16 vs 8）能获得更高的接受长度，尤其是在数学和代码任务上（见表7）。一个关键的发现是，使用较大块（如16）训练的模型能很好地泛化到推理时使用较小的块（如8），但反之则不然。这一特性为推理时动态调度块大小以优化端到端效率提供了可能。

表7. 训练-推理块大小(BS)不匹配的消融研究。
图片描述

A5 结论

本文提出了DFlash，一个基于扩散的推测解码框架，它重新思考了扩散语言模型在加速自回归LLM推理中的作用。通过将扩散模型的功能限制在草稿生成阶段，DFlash利用了其固有的并行性，同时避免了限制其独立使用的质量下降问题。通过将扩散草稿模型置于丰富的目标模型上下文条件下，DFlash实现了高接受率，从而将推理速度显著提升至超越先前推测解码方法的水平。

除了经验上的改进，DFlash还为扩散式LLM提出了一种新的发展范式。扩散模型可以不必在端到端生成质量上与自回归模型竞争，而是可以作为轻量级、专门的草稿生成器，为快速准确的块预测进行优化。这种重新定位允许对去噪步骤进行激进的削减以最大化并行性，而推测性验证则为输出质量提供了原则性的保证。我们希望DFlash能将基于扩散的草稿生成确立为一种实用且可扩展的推测解码范式，推动更高效、更易于访问的LLM部署。

A6 附录

A.1. 训练实现

优化器和学习率。草稿模型使用AdamW优化器训练6个epoch，学习率为$6 \times 10^{-4}$，梯度裁剪阈值为1.0，并采用带有0.04预热比例的余弦学习率调度。

训练配置。我们在训练数据混合集上进行训练，最大序列长度为3072个token（Qwen3-Coder为4096）；对于每个序列，随机采样512个锚点位置。公式4中损失衰减的超参数$\gamma$对于块大小为16的模型设置为7，块大小为10的模型设置为5，块大小为8的模型设置为4。

训练模式。训练可以采用在线或离线方式进行。在线训练中，目标模型的隐藏特征在每个训练步骤中动态计算。离线训练中，目标模型的隐藏特征被预先计算并缓存，然后在草稿模型优化期间加载，以减少计算开销。

A.2. 不使用目标特征的扩散草稿模型

实验结果。表8展示了一个不使用目标上下文特征的5层块扩散草稿模型。结果显示，该模型仅取得了中等的接受长度和加速比，这验证了利用目标模型特征进行条件化对于实现高性能至关重要。

表8. 一个不含目标上下文特征的5层块扩散草稿模型。该草稿模型仅实现了中等的接受长度和加速比。
图片描述

A.3. 进一步的消融研究

A.3.1. 损失衰减

对比实验。我们对4.2节中介绍的与位置相关的损失衰减进行了消融研究。具体来说，我们将使用指数衰减token权重的默认设置与在每个草稿块内使用均匀token权重训练的变体进行了比较。这项研究旨在分离出在训练期间强调早期token准确性的效果。

结论。图5中的结果显示，应用损失衰减可以使训练收敛得更快、效果更好。

图5. 损失衰减使训练收敛得更快、效果更好。
图5. 损失衰减使训练收敛得更快、效果更好。

A.3.2. 掩码块的随机采样

对比实验。我们比较了两种构建掩码块的策略：标准方法（将响应均匀划分为块）和我们的方法（随机采样锚点token来构建块）。两个3层草稿模型都在10万样本数据上训练，并从目标模型中提取5个隐藏特征，块大小为16。

结论。如表9所示，在训练期间随机采样锚点token来构建掩码块，有效地增强了训练数据，并带来了更高的接受长度和更好的加速比。

表9. 训练期间随机采样锚点token来构建掩码块能有效增强训练数据，并带来更高的接受长度和更好的加速比。两个草稿模型都使用三层，并从目标模型中提取五个隐藏特征。块大小为16。我们使用第5.4节介绍的10万样本数据来训练这两个模型。
图片描述