DFlash:基于块扩散模型的快速推测解码框架 DFlash: Block Diffusion for Flash Speculative Decoding
用轻量块扩散模型做并行草稿,实现6倍以上无损加速
前置知识
自回归语言模型(Autoregressive LLM)
自回归语言模型是当前主流的文本生成范式,每次生成一个 token,且每个 token 的生成都依赖于前面所有 token 的上下文。这种串行生成方式意味着生成 $n$ 个 token 需要 $n$ 次前向传播,导致推理延迟高、GPU 利用率低。GPT、LLaMA、Qwen 等主流大模型均采用这一范式。
DFlash 的核心目标就是解决自回归解码的串行瓶颈,理解这一范式的固有限制是理解本文动机的前提。
推测解码(Speculative Decoding)
推测解码是一种加速 LLM 推理的范式:用一个轻量级的草稿模型(draft model)快速生成若干候选 token,然后由目标大模型(target model)并行验证这些 token。如果草稿 token 与目标模型的分布一致则接受,否则从拒绝位置重新采样。这种方法能实现无损加速——输出分布与目标模型完全一致。代表性方法包括 EAGLE 系列、Medusa 等。
DFlash 是推测解码的一种新变体,用扩散模型替代自回归草稿模型,理解推测解码的基本机制是理解本文创新的基础。
扩散语言模型(Diffusion Language Model, dLLM)
扩散语言模型借鉴了图像扩散生成的思想,通过迭代去噪过程生成文本。与自回归模型不同,扩散模型可以在每一步并行预测多个 token。块扩散模型(Block Diffusion)是其中一种变体,将序列分成若干块,每块内的 token 并行去噪。代表性工作包括 LLaDA、Fast-dLLM v2、SDAR 等。然而,现有 dLLM 的生成质量通常不及同等规模的自回归模型,且需要较多去噪步骤。
DFlash 利用块扩散模型作为草稿模型实现并行生成,理解扩散模型的工作原理和优缺点是理解本文技术路线的关键。
KV 缓存(Key-Value Cache)
在 Transformer 模型的推理过程中,KV 缓存存储了每一层 Attention 中的 Key 和 Value 向量,避免在生成新 token 时重复计算之前 token 的表示。这使得每步推理只需要计算当前 token 的 Query,大幅降低计算量。KV 缓存是 Transformer 推理优化的核心组件,其大小直接影响内存占用和推理效率。
DFlash 的核心创新之一是将目标模型的隐藏特征注入草稿模型的 KV 缓存(而非仅作为输入),这一设计直接决定了草稿质量能否随草稿模型深度有效提升。
接受长度(Acceptance Length, $\tau$)
接受长度是推测解码中的核心指标,表示每个解码周期中平均被目标模型接受的 token 数量(包括目标模型额外生成的一个 bonus token)。接受长度 $\tau$ 的取值范围在 $[1, \gamma+1]$ 之间,其中 $\gamma$ 是草稿 token 数。接受长度越长,说明草稿质量越高,加速效果越好。
DFlash 的核心优势在于通过扩散并行生成和目标模型特征注入实现了远高于 EAGLE-3 的接受长度(6.5-7.9 vs 3.0-3.7),这是理解其加速效果的关键。
EAGLE 系列方法
EAGLE 是当前最先进的推测解码方法系列。EAGLE-1 通过预测目标模型的未来隐藏状态分布来提升接受率;EAGLE-2 引入自适应草稿树;EAGLE-3 进一步改进训练目标并使用单层 Transformer 作为草稿模型。EAGLE 系列利用目标模型的特征级上下文信息,但仍然采用自回归方式生成草稿 token,因此草稿延迟随 token 数线性增长。
EAGLE-3 是本文的主要基线方法,理解其工作原理和局限性有助于理解 DFlash 的改进之处——从自回归草稿转向扩散并行草稿。
研究动机
自回归大语言模型的推理存在严重的串行瓶颈:每个输出 token 的生成都依赖完整的前置上下文,导致推理延迟高、GPU 利用率低。随着长链式推理(Chain-of-Thought)模型的兴起,这一问题变得更加突出,因为推理时间现在主导了整个生成过程。推测解码是当前的主要解决方案,但最先进的方法如 EAGLE-3 仍然依赖自回归草稿生成——草稿 token 是逐个串行生成的。这种串行草稿过程不仅本质上低效,还容易产生错误累积,将可实现的加速比限制在大约 2-3 倍。具体来看,EAGLE-3 使用单层 Transformer 作为草稿模型,在 Qwen3-8B 上平均加速仅为 2.02 倍(树大小 60),接受长度仅为 3.40。另一方面,扩散语言模型提供了并行生成的可能性,但现有的开源 dLLM 通常在生成质量上不及自回归模型,且需要大量去噪步骤,严重影响了实际推理速度。
本文的目标是本文的目标是构建一个既轻量又高精度的扩散草稿模型,充分利用推测解码框架的优势,实现远超现有方法的无损加速。具体而言,DFlash 希望:(1) 利用块扩散模型实现草稿 token 的单次前向传播并行生成,大幅降低草稿延迟;(2) 通过从目标模型提取的隐藏特征对草稿模型进行条件化,提高草稿质量和接受长度;(3) 在多种模型和任务上验证方法的通用性,实现在真实服务场景下的实际加速。
与已有工作不同的是,本文的独特切入角度在于一个关键洞察:「目标模型最了解未来」。大型自回归 LLM 的隐藏特征隐式地编码了关于多个未来 token 的信息。DFlash 利用这些隐藏特征作为上下文,条件化草稿模型以并行预测未来的 token 块。本质上,草稿模型变成了一个扩散适配器,高效利用大目标模型建模的深层上下文信息。与之前的工作不同,DFlash 不是让一个小草稿模型从头推理,而是将目标模型的推理能力与小扩散草稿器的并行生成速度相融合。此外,DFlash 采用了全新的特征注入方式——将目标特征注入草稿模型每一层的 Key 和 Value 投影中(而非仅在输入层融合),这使得接受长度能随草稿模型深度有效增长,突破了现有方法的加速瓶颈。
核心方法
DFlash 的整体思路可以概括为「用扩散做草稿,用目标模型做验证」。直觉上,自回归模型在生成下一个 token 时,其内部隐藏状态实际上已经包含了关于未来多个 token 的部分信息——就像一个经验丰富的作者在写一个词时,脑海中已经有了接下来几个词的轮廓。DFlash 的做法是:首先让目标大模型完成一次标准的 prefill 过程并提取其隐藏特征;然后用一个轻量的块扩散模型,以这些隐藏特征为上下文,在单次前向传播中并行生成一批草稿 token;最后由目标模型并行验证这些草稿 token。技术路线上,DFlash 包含三个关键组件:(1) 目标上下文特征提取与融合——从目标模型的多个层提取隐藏状态并通过投影层融合;(2) KV 注入条件化——将融合后的特征注入草稿模型每一层的 Key 和 Value 中;(3) 块扩散并行草稿——使用块扩散过程在单次前向传播中生成多个草稿 token。
DFlash 的核心创新在于两个方面,与已有方法有本质区别。第一个创新是 KV 注入(KV Injection)替代输入融合(Input Fusion)。EAGLE-3 等方法也将目标模型的隐藏特征与草稿模型的 token 嵌入融合,但只在输入层注入。随着草稿模型深度增加,目标信息会逐渐被稀释,导致增加草稿层数时接受长度收益递减。DFlash 将融合后的目标上下文特征视为持久的上下文信息,直接注入每一层草稿模型的 Key 和 Value 投影中,存储在草稿模型的 KV 缓存中并跨草稿迭代复用。这种设计在整个草稿模型中提供了强而一致的条件化,使接受长度能随草稿层数有效增长。第二个创新是将块扩散模型定位为轻量级草稿适配器。与 DiffuSpec、SpecDiff-2 等方法使用巨大的 7B 参数扩散模型作为草稿器不同,DFlash 仅使用 5 层 Transformer 作为草稿模型,共享目标模型的 embedding 和 LM head,仅训练草稿 Transformer 层。这种设计使草稿模型成为一个高效利用目标模型深层信息的扩散适配器。
方法步骤详情
DFlash 的推理流程包含以下步骤:(1) 目标模型预填充:给定输入提示,目标模型首先执行标准的 prefill 过程生成第一个 token;(2) 隐藏特征提取:在 prefill 过程中,从目标模型的固定层集合(均匀采样从浅层到深层)提取隐藏表示,共提取 5 层特征;(3) 特征融合:将提取的隐藏状态拼接后通过一个轻量投影层,融合跨层信息为紧凑的目标上下文特征;(4) KV 缓存注入:将融合后的目标上下文特征注入草稿模型每一层的 Key 和 Value 投影中;(5) 块扩散草稿生成:草稿模型使用块扩散过程,在单次前向传播中并行生成 $\gamma$ 个草稿 token(默认 $\gamma=16$);(6) 目标模型并行验证:目标模型并行验证所有草稿 token,接受与目标分布一致的 token,拒绝位置重新采样;(7) 迭代:接受的 token 中最后一个作为下一轮的「bonus token」(锚点),重复步骤 2-6。训练方面,DFlash 有几个关键设计:随机锚点采样——从响应中随机采样锚点 token 作为块的起始位置,掩码剩余位置,直接匹配推理时的行为;损失加权——对块内早期位置的交叉熵损失施加指数衰减权重 $w_k = \exp(-k/\gamma)$,因为早期位置的错误会使后续所有 token 无效;共享嵌入和 LM head——草稿模型共享目标模型的 token embedding 层和语言模型头并保持冻结,仅训练草稿 Transformer 层。
技术新颖性
DFlash 的技术新颖性体现在以下几个层面。首先,它首次将扩散模型定位为推测解码框架中的轻量级草稿适配器,而非独立的生成模型,这一视角转换使得扩散模型可以规避生成质量不足的问题,同时充分发挥其并行生成的优势。其次,KV 注入机制是一种全新的特征条件化方式,与 EAGLE 系列的输入融合有本质区别:输入融合只在第一层注入目标信息,信息会随深度稀释;而 KV 注入在每一层都维持目标信息的强条件化,使接受长度能随草稿深度有效增长。实验表明,KV 注入在自回归草稿模式下也优于输入融合(接受长度 4.8 vs 4.2),在块扩散模式下优势更加明显(4.2 vs 3.5)。第三,DFlash 的训练策略针对推测解码场景进行了专门设计,包括随机锚点采样(匹配推理时的锚点条件化行为)、损失加权(强调早期位置以提高接受率)、以及稀疏注意力掩码(块内双向注意力、块间不可见)。这些设计共同使得一个仅有 5 层的轻量草稿模型能够实现远超 EAGLE-3 的加速效果。
实验结果
DFlash 在多个模型和任务上实现了显著的加速效果。在 Qwen3-8B 上使用 Transformers 后端的实验中(温度 0,最大生成 2048 token),DFlash 在 Math500 上实现了 6.08 倍加速(接受长度 7.87),而 EAGLE-3 在树大小 16 时仅为 1.81 倍(接受长度 3.02),在树大小 60 时为 2.05 倍(接受长度 3.49)。在 GSM8K 上,DFlash 实现 5.15 倍加速,EAGLE-3 (16) 为 1.94 倍,EAGLE-3 (60) 为 2.23 倍。在代码任务上,HumanEval 上 DFlash 达到 5.14 倍加速,MBPP 为 4.65 倍,LiveCodeBench 为 5.51 倍。平均而言,DFlash 在 Qwen3-8B 上实现 4.86 倍平均加速,是 EAGLE-3 (16) 的 2.75 倍,是 EAGLE-3 (60) 的 2.41 倍。在推理模型(thinking mode 开启)场景下,DFlash 在 Qwen3-4B 的 MATH-500 上实现 4.59 倍加速,Qwen3-8B 上为 4.64 倍,证明了其在长推理场景下的有效性。在 SGLang 服务框架上的真实部署实验中(单块 B200 GPU),DFlash 在 Qwen3-8B Math500 上并发 1 时实现 5.1 倍加速,并发 32 时仍有 2.8 倍加速。对于更大的 Qwen3-Coder-30B-A3B 模型,DFlash 在 HumanEval 上并发 1 时实现 3.5 倍加速,并发 32 时为 3.1 倍,展示了良好的可扩展性。在 LLaMA-3.1-8B-Instruct 上,DFlash 在 GSM8K 上并发 1 时实现 2.4 倍加速,HumanEval 上为 2.8 倍,均超过 EAGLE-3。长上下文适应实验表明,通过仅 1600 个样本的微调,DFlash 草稿模型可以从 4K 上下文扩展到 32K,接受长度从 2.09 提升到 3.56。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GSM8K(数学推理) | 加速比 / 接受长度 $\tau$ | 5.15× / 6.54 | EAGLE-3 (16): 1.94× / 3.23; EAGLE-3 (60): 2.23× / 3.71 | 相对 EAGLE-3 (16) 提升 2.65 倍,相对 EAGLE-3 (60) 提升 2.31 倍 |
| MATH-500(数学推理) | 加速比 / 接受长度 $\tau$ | 6.08× / 7.87 | EAGLE-3 (16): 1.81× / 3.02; EAGLE-3 (60): 2.05× / 3.49 | 相对 EAGLE-3 (16) 提升 3.36 倍,相对 EAGLE-3 (60) 提升 2.97 倍 |
| AIME25(竞赛数学) | 加速比 / 接受长度 $\tau$ | 5.62× / 7.08 | EAGLE-3 (16): 1.79× / 3.00; EAGLE-3 (60): 2.05× / 3.44 | 相对 EAGLE-3 (16) 提升 3.14 倍,相对 EAGLE-3 (60) 提升 2.74 倍 |
| HumanEval(代码生成) | 加速比 / 接受长度 $\tau$ | 5.14× / 6.50 | EAGLE-3 (16): 1.89× / 3.17; EAGLE-3 (60): 2.17× / 3.65 | 相对 EAGLE-3 (16) 提升 2.72 倍,相对 EAGLE-3 (60) 提升 2.37 倍 |
| MBPP(代码生成) | 加速比 / 接受长度 $\tau$ | 4.65× / 5.95 | EAGLE-3 (16): 1.69× / 2.82; EAGLE-3 (60): 1.93× / 3.25 | 相对 EAGLE-3 (16) 提升 2.75 倍,相对 EAGLE-3 (60) 提升 2.41 倍 |
| LiveCodeBench(代码推理) | 加速比 / 接受长度 $\tau$ | 5.51× / 7.27 | EAGLE-3 (16): 1.57× / 2.65; EAGLE-3 (60): 1.81× / 3.03 | 相对 EAGLE-3 (16) 提升 3.51 倍,相对 EAGLE-3 (60) 提升 3.05 倍 |
| MT-Bench(对话质量) | 加速比 / 接受长度 $\tau$ | 2.75× / 4.24 | EAGLE-3 (16): 1.63× / 2.83; EAGLE-3 (60): 1.90× / 3.26 | 相对 EAGLE-3 (16) 提升 1.69 倍,相对 EAGLE-3 (60) 提升 1.45 倍 |
| SGLang Math500(Qwen3-8B,并发1) | 吞吐量加速比 | 5.1×(1175 tok/s vs 基线 230 tok/s) | EAGLE-3 (10): 1.6×; EAGLE-3 (60): 未报告 | 5.1 倍吞吐量提升 |
局限与改进
尽管 DFlash 取得了显著的加速效果,但仍存在一些局限性。首先,DFlash 的加速效果在对话任务(MT-Bench)上相对较弱,Qwen3-8B 上仅为 2.75 倍,而在数学和代码任务上可达 5-6 倍。这可能是因为对话任务的生成模式更加多样和不可预测,扩散模型的并行预测在这种场景下优势较小。其次,DFlash 需要为目标模型单独训练草稿模型,增加了部署的复杂性。虽然训练成本相对较低(使用约 800K 样本),但针对不同的目标模型需要分别训练。第三,论文未与其他基于扩散的推测解码方法(如 DiffuSpec、SpecDiff-2、TiDAR 等)进行直接比较,原因是缺乏开源实现,这使得 DFlash 的相对优势难以完全量化。第四,DFlash 在高并发场景下(如并发 32)的加速比会下降,Qwen3-8B Math500 从并发 1 的 5.1 倍降至并发 32 的 2.8 倍,这反映了推测解码在计算受限场景下的固有局限。此外,作者也指出在大批次场景下,较大的块大小可能增加验证成本,需要动态块大小调度来优化。
独立分析的弱点
DFlash 存在几个值得改进的弱点。第一,对话任务加速不足:在 MT-Bench 上的加速比仅为 2.75 倍,远低于数学和代码任务的 5-6 倍,可能需要针对对话场景优化训练数据或调整块大小策略。第二,块大小固定:当前方法使用固定的块大小(16 for Qwen3,10 for LLaMA-3.1),无法根据生成内容的可预测性动态调整。作者在论文中也提到,大批次下大块会增加验证成本,自适应块大小调度是一个有价值的改进方向。第三,训练数据依赖:DFlash 使用目标模型生成的响应作为训练数据以实现更好的目标对齐,这意味着需要先用目标模型生成大量数据,增加了前期成本。第四,特征提取开销:推理时需要从目标模型的多层提取隐藏特征并进行投影融合,虽然开销不大,但在极低延迟场景下可能成为瓶颈。第五,温度敏感性:在温度 1(非贪婪采样)设置下,DFlash 的加速效果有所下降(Qwen3-8B 平均从 4.86 倍降至 4.03 倍),而 EAGLE-3 的下降幅度相对较小,说明扩散草稿对采样温度更敏感。
未来方向
论文和基于当前成果可延伸出多个研究方向。作者提出的直接方向包括:(1) 自适应块大小调度——根据当前的计算受限或内存受限状态动态调整草稿块大小,在大批次场景下使用小块降低验证成本,在小批次场景下使用大块提高并行度;(2) 长上下文扩展——论文已展示了通过少量样本微调将草稿模型从 4K 扩展到 32K 的可行性,未来可进一步探索更长上下文的支持。基于当前成果可延伸的方向包括:(3) 多模态推测解码——将 DFlash 的扩散草稿思想扩展到视觉语言模型等多模态场景;(4) 草稿模型蒸馏——探索更小的草稿模型(如 2-3 层)是否能在保持加速效果的同时进一步降低草稿延迟;(5) 与投机采样结合——探索 DFlash 与投机采样(Speculative Sampling)在可控生成、对齐等场景下的结合;(6) 动态特征层数选择——根据任务类型自适应选择从目标模型提取多少层特征,平衡特征质量和计算开销。
复现评估
论文的复现条件较好。代码已在 GitHub 开源(链接见论文),模型权重在 Hugging Face 发布。训练数据方面,DFlash 使用约 800K 样本,来自 NVIDIA Nemotron Post-Training Dataset V2 和 CodeAlpaca,并使用目标模型重新生成响应。算力要求方面,所有实验在 NVIDIA H200 GPU 上进行,SGLang 部署实验使用单块 B200 GPU。训练细节方面,草稿模型使用 5 层 Transformer(Qwen3 Coder 为 8 层),块大小 16(LLaMA-3.1 为 10),目标隐藏特征从 5 层均匀提取。论文提供了详细的实现细节,包括损失加权公式、注意力掩码设计、训练数据构造方法等。整体而言,复现难度中等——核心算法清晰,但需要较多 GPU 资源进行训练和评估,且需要对目标模型进行推理以生成训练数据。
论文图表