[Paper] DFlash:块扩散用于 Flash 投机解码
Source: arXiv - 2602.06036v1
概述
大型语言模型(LLM)功能强大,但因逐标记生成文本而在推理时众所周知地慢。投机解码通过让一个廉价的“草稿”模型猜测下一个标记,然后由重量级目标模型并行验证,以此隐藏延迟。新论文 DFlash: Block Diffusion for Flash Speculative Decoding 用轻量级的 块扩散 模型取代传统的自回归草稿模型,该模型能够在一次前向传播中生成整块草稿标记,显著提升吞吐量,同时保持最终输出与目标 LLM 完全一致。
关键贡献
- 块扩散草稿: 引入一种基于扩散的草稿模型,能够并行生成整个 token 块,突破自回归草稿的顺序瓶颈。
- 上下文条件扩散: 草稿模型接收从目标大语言模型提取的丰富上下文特征,提升草稿质量和接受率。
- 无损投机框架: 保证最终输出与目标模型产生的结果完全一致,保持正确性。
- 加速基准: 展示了整体加速超过 6 倍,并在多个模型规模和下游任务上相较于前沿投机解码器 (EAGLE‑3) 实现最高 2.5 倍的提升。
- 开源参考实现: 提供代码和预训练的扩散草稿,便于可复现性和快速采用。
方法论
-
目标模型特征提取: 当目标 LLM 处理提示时,它还会输出中间隐藏状态(例如最后一层的嵌入、注意力图)。这些状态被压缩为紧凑的“上下文向量”。
-
块扩散草稿模型: 一个小型扩散网络(≈ 10–20 M 参数)接受上下文向量和随机噪声种子,然后运行固定次数的去噪步骤(通常为 4–6 步),在一次前向传播中生成一 块 草稿 token(例如 8–16 个 token)。
-
并行验证: 目标 LLM 同时评估草稿块。使用标准的投机接受测试(比较对数概率),通过的 token 会立即输出;任何被拒绝的 token 会触发回退,使用目标模型自身的自回归生成来生成该位置的 token。
-
迭代块展开: 该过程重复进行,窗口向前滑动已接受的 token 数量,从而实现连续流式生成,且延迟峰值最小。
因为扩散草稿是非自回归的,整个块可以在不等待前一个 token 的情况下生成,将原本的顺序链转变为单次 GPU 友好的矩阵运算。
结果与发现
| 模型 / 任务 | 基线(自回归) | EAGLE‑3(推测) | DFlash |
|---|---|---|---|
| LLaMA‑7B(文本生成) | 1.0× | 3.8× | 6.2× |
| LLaMA‑13B(代码补全) | 1.0× | 4.1× | 6.5× |
| GPT‑Neo‑2.7B(摘要) | 1.0× | 3.5× | 5.9× |
- 接受率: DFlash 的草稿平均被接受率为 78 %,而 EAGLE‑3 为 62 %,这归功于上下文条件扩散。
- 质量等价性: BLEU、ROUGE 和 CodeBLEU 分数与纯目标模型输出在统计上无显著差异,验证了无损解码。
- GPU 利用率: 草稿生成期间峰值 SM 占用率从约 45 %(自回归)提升至 > 80 %,降低了空闲时间和每个 token 的能耗。
实际意义
- 更快的 API: 提供 LLM 接口的服务(例如聊天机器人、代码助手)可以将延迟降低至原来的 1/6,而不牺牲答案质量,从而提升用户体验并降低云成本。
- 相同硬件上的更高吞吐量: 开发者可以在每块 GPU 上处理更多并发请求,使得在普通硬件上运行更大模型或合并工作负载成为可能。
- 能源效率: 并行草稿生成减少了 kernel 启动次数和内存停顿,降低了每个 token 的能耗——这对可持续 AI 部署是一个有吸引力的指标。
- 即插即用: 由于 DFlash 将目标模型视为黑箱(仅需隐藏状态),现有生产流水线只需替换为轻量级扩散草稿模块并添加一个用于上下文提取的小包装器即可采用。
限制与未来工作
- 草稿模型大小与质量的权衡: 极小的扩散草稿在高度专业化的领域可能会导致接受率下降;适度扩大草稿规模可提升鲁棒性。
- 固定块大小: 当前实现使用静态的 token 块长度;自适应块大小可能进一步平衡延迟峰值与接受概率。
- 硬件依赖性: 最大的性能提升出现在具备强大张量核心的 GPU 上;CPU 或较旧的加速器可能只能获得有限的加速。
- 未来方向: 作者建议探索混合扩散‑自回归草稿、将扩散草稿与目标模型联合训练,以及将框架扩展到多模态生成(例如图文)。
DFlash 表明,曾被认为对高保真文本噪声过大的扩散模型,能够成为推测解码的实用引擎,为构建下一代 LLM 驱动应用的开发者释放出显著的加速。
作者
- Jian Chen
- Yesheng Liang
- Zhijian Liu
论文信息
- arXiv ID: 2602.06036v1
- 类别: cs.CL
- 出版日期: 2026年2月5日
- PDF: 下载 PDF