DFlash

DFlash 是面向投机解码的块并行 Draft Model。它用一次 Forward 预测一组候选 token，并把 Target Model 的多层隐藏状态持续注入 Draft Model 每一层的 KV Cache；Target Model 仍按标准规则验证候选，因此加速不会改变目标分布。

本文解释的是论文 DFlash: Block Diffusion for Flash Speculative Decoding 中的 DFlash。官方仓库后来加入了 DFlash 2，两者的模型结构与实验结果不应混用。

1. 动机：降低串行 Draft 成本

设一轮提出 $\gamma$ 个候选，Draft 与 Target 验证耗时分别为 $T_{\text{draft}}$、$T_{\text{verify}}$，该轮平均产出 $\tau$ 个 token，则平均单 token 延迟为：

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

自回归 Draft Model 必须逐 token 生成候选。若单步耗时为 $t_d$，则近似有：

$$T_{\text{draft}}^{\text{AR}}\approx \gamma t_d$$

这会形成直接冲突：增加 Draft Model 深度或候选长度通常能提高 $\tau$，但也会线性增加串行步骤。DFlash 将候选位置放在同一个块中并行计算，使 Draft 阶段只需一次模型 Forward。实际耗时仍会随块大小、注意力后端和硬件变化，但不再包含 $\gamma$ 次模型级串行调用。

自回归 Draft：x₁ → x₂ → ... → xγ       γ 次串行 Forward
DFlash：      [x₁, x₂, ..., xγ]         1 次块并行 Forward
Target：      [x₁, x₂, ..., xγ]         1 次并行验证

因此 DFlash 同时优化两个量：块并行降低 $T_{\text{draft}}$，KV 注入提高候选质量与 $\tau$。

2. 模型结构

2.1 Target Context Feature

DFlash 不让小型 Draft Model 独立重建长上下文语义，而是从 Target Model 的若干层抽取隐藏状态。设选中的层为 $l_1,\dots ,l_m$，拼接后通过共享投影与 RMSNorm 得到上下文特征：

$$H_t=\mathrm{RMSNorm}(W_c[H^{(l_1)};\dots ;H^{(l_m)}])$$

论文的主实验通常均匀选择 5 个 Target 层，并使用 5 层 Draft Model；这只是具体配置，不是算法约束。抽取更多层会增加特征缓存与训练成本，Draft Model 更深也会增加单次 Forward 延迟。

2.2 KV 注入

输入融合只在 Draft Model 的入口注入 Target 特征，信号可能随网络深度被稀释。DFlash 改为在每个 Draft Layer 中把 $H_t$ 作为额外的 Key、Value：

$$Q_i=W_i^Q{H}_d$$$$K_i=[W_i^K{H}_t;W_i^K{H}_d{]}_{\text{seq}},V_i=[W_i^V{H}_t;W_i^V{H}_d{]}_{\text{seq}}$$

其中 $H_d$ 是当前候选块的隐藏状态。Target 特征不经过 Draft Layer 的 Query、Attention 输出和 FFN，只作为持久上下文写入 Draft Model 自己的 KV Cache。它并不是复用 Target Model 原有的 K、V 张量，而是用 Target 隐藏状态重新做 Draft 侧投影。

这种设计让每一层、每一轮 Draft 都能访问同一份高质量上下文。论文消融中，无论自回归 Draft 还是块并行 Draft，KV 注入都比只在输入端融合获得更高的平均接受长度。

2.3 单次块并行预测

DFlash 训练时从回复中选择一个干净的 anchor token，把块内其余位置替换为 mask token。块内采用非因果注意力，使所有 mask 位置在一次 Forward 中相互可见并并行输出 logits。

这里的“块扩散”描述训练掩码与块级并行生成方式。原始 DFlash 推理不是通用 Diffusion LLM 的多轮去噪：它没有反复 Mask、Remask 或迭代收敛，而是一次 Forward 直接产生整块候选。最终正确性来自 Target Model 的投机验证，不依赖 Draft 本身生成最终文本。

论文与推理引擎的块大小口径还可能不同：论文中的 block size 包含 anchor，当前 vLLM 的 num_speculative_tokens=N 表示 $N$ 个候选，并构造 1 + N 个查询位置。比较配置时应先统一“是否包含 anchor/bonus token”。

3. 训练方法

3.1 随机 Anchor 与稀疏注意力

训练序列先经过冻结的 Target Model，得到所有 token 的多层上下文特征；随后随机采样多个 anchor，并为每个 anchor 构造独立的 mask block。多个块拼接后在一次 Forward/Backward 中训练，但注意力掩码必须满足：

• 块内 token 可以双向关注同一块；
• 每个块只能访问与自身位置对应的 Target Context Feature；
• 不同块之间不可互相泄露未来 token。

随机 anchor 既匹配“上一轮 Target 产出的 bonus token 作为下一轮起点”的推理流程，也比固定切块覆盖更多上下文位置。

3.2 位置衰减损失

投机解码只接受连续前缀，前部错误会使整个后缀失效。DFlash 因此对块内第 $k$ 个预测位置使用指数衰减权重：

$$w_k=\exp (-\frac{k-1}{\lambda })$$$${\mathcal{L}}_{\text{DFlash}}=-\sum _{k=1}^{\gamma }{w}_k\log q_k(x_k^{\ast })$$

这里用 $\lambda$ 表示损失衰减系数，以免与候选长度 $\gamma$ 混淆。Target Model、共享 Token Embedding 和 LM Head 均保持冻结，只更新 Draft Transformer 与上下文投影。

4. 推理与分布等价性

一轮 DFlash 推理包含以下步骤：

1. Target Model 对新上下文计算 logits 和选定层的隐藏状态，并产生 anchor/bonus token。
2. Draft Model 将 Target 隐藏状态投影为各层 K、V，写入 Draft KV Cache。
3. Draft Model 对 mask block 做一次并行 Forward，得到 $\gamma$ 个候选分布 $q_1,\dots ,q_{\gamma }$。
4. Target Model 一次 Forward 得到验证分布 $p_1,\dots ,p_{\gamma +1}$。
5. 贪心模式接受最长匹配前缀；采样模式按标准拒绝采样逐位置校正。

采样模式下，候选 $x_k$ 的接受概率仍是：

$$A(x_k)=min(1,\frac{p_k(x_k)}{q_k(x_k)})$$

拒绝后从归一化残差 $max(0,p_k-q_k)$ 采样。DFlash 只改变提议分布 $q$ 的计算方式，不改变校正规则；只要后端保留 Draft 概率并正确执行验证，输出分布就与 Target Model 一致。仅比较 token 是否相同的贪心验证不能替代温度采样下的概率校正。

5. vLLM 实现

截至 2026-08-19，vLLM 的 DFlash 路径把实现拆成模型与 Speculator 两部分：

组件	职责
qwen3_dflash.py	融合 Target 辅助隐藏状态，批量投影各 Draft Layer 的 K、V，并写入 Draft KV Cache
dflash/speculator.py	构造 anchor + mask 查询块、执行并行 Forward、采样候选并维护缓存位置
验证与采样器	复用 vLLM 的投机验证逻辑，支持贪心或概率拒绝采样

KV 物化路径会先把所有 Draft Layer 的 K/V 权重拼接，再用一次 GEMM 完成跨层投影，随后统一处理 K-Norm、RoPE 和 KV Cache 写入。这样避免逐层启动小型投影 kernel。

vLLM：跨层融合的 Context KV 投影

normed_context_states = torch.empty_like(context_states)
ops.rms_norm(
    normed_context_states,
    context_states,
    self._hidden_norm_weight,
    self._rms_norm_eps,
)
all_kv_flat = F.linear(
    normed_context_states,
    self._fused_kv_weight,
    self._fused_kv_bias,
)
all_kv = all_kv_flat.view(
    num_ctx, num_layers, 2, num_kv_heads, head_dim
).permute(2, 1, 0, 3, 4).contiguous()

源码：vLLM qwen3_dflash.py

当前实现还存在明确边界：DFlash Speculator 不支持多模态输入；非因果 Draft Attention 需要对应后端支持；混合 Sliding/Full Attention 的检查点需要 V2 Model Runner。模型配置、Draft checkpoint、Target RoPE 布局和 mask token 也必须一致，否则常见现象不是程序报错，而是接受率显著下降。

6. 实验结果如何解读

论文 Table 1 在单张 NVIDIA H200、Transformers 后端、最大生成 2048 token、Qwen3 thinking 关闭的条件下报告：

Target	设置	EAGLE-3（预算 16）	DFlash（block 16）
Qwen3-4B	greedy，7 个任务平均	1.81× / $\tau =3.05$	4.91× / $\tau =6.54$
Qwen3-8B	greedy，7 个任务平均	1.76× / $\tau =2.96$	4.86× / $\tau =6.49$
Qwen3-4B	temperature 1	1.72× / $\tau =2.95$	4.24× / $\tau =5.69$
Qwen3-8B	temperature 1	1.68× / $\tau =2.83$	4.03× / $\tau =5.48$

表中前一个数是相对原生自回归的解码加速，后一个数是含 bonus token 的平均接受长度。它说明块并行在该实验设置下同时降低了 Draft 延迟并提高了接受长度，但不能直接外推到任意模型和 Serving 负载。

论文的 SGLang 实验进一步展示了并发影响：Qwen3-8B、Math500、单张 B200 + FA4 下，并发 1 时为 5.1×，并发 32 时降至 2.8×。原因是 Batch 增大后 Target 验证逐渐转为计算受限，额外验证低概率后缀不再接近免费。

7. 适用边界

• 需要 Target 专属训练。 Draft 依赖特定 Target 的层选择、隐藏维度、词表和表示空间，不能像通用小模型那样随意替换。
• 长上下文不会自动泛化。 论文中 4K 训练的基础模型在更长上下文上接受长度下降，需要少量长上下文数据继续训练。
• 独立位置存在后缀衰减。 每个候选位置不知道块内前一个位置最终采样了什么，长块容易发生多模态碰撞；DSpark 用轻量顺序头处理这一问题。
• 最佳块大小依赖负载。 大块在低并发下能提高每轮产出，高并发下可能增加无效验证。论文也把自适应块大小留作后续工作。
• “无损”依赖完整验证链路。 Draft Model 本身不保证质量；只有正确的 Target 验证与拒绝采样才保证分布等价。

8. 参考资料

1. DFlash: Block Diffusion for Flash Speculative Decoding
2. DFlash 官方仓库
3. DeepSpec：DFlash / DSpark / EAGLE-3 训练与评测
4. vLLM DFlash Speculator
5. vLLM Qwen3 DFlash Model