核心思想

将Diffusion LLM的思想应用到草稿模型上,具体实现上,为了KV cache,采用Block Diffusion的方案实现。

模型架构

DFlash

DFlash

Eagle3

Eagle3

不同于Eagle3(绿色标记),DFlash(红色标记)将融合的目标上下文特征视为持久性上下文信息,并将其直接注入到每个草稿模型层的 Key 和 Value 投影中。投影后的特征存储在草稿模型的 KV 缓存中,并在多次草稿生成迭代中重复使用。这种设计在整个草稿模型中提供了强大且一致的条件约束,使可接受长度能够随着草稿层数量的增加而有效扩展。

DFlash还有一点和Eagle3不同的是,DFlash的草稿模型有多层Decoder Layer,而Eagle3仅一层。

DFlash的关键是其掩码注意力的实现。

The draft model is trained to predict the next block_size - 1 tokens in parallel.

蓝色(Target Context Feature):训练时,先把 “输入提示词 + 正确回复” 的完整序列传给目标模型,从它的中间层提取隐藏特征(包含语义逻辑、推理轨迹、未来 token 的潜在信息),再通过融合得到 “目标上下文特征”—— 相当于目标大模型把 “想表达的思路” 提炼出来,告诉草稿模型。

黄色(Clean Token):“正确回复” 中随机挑选的 “锚点 token”,是每个草稿块的起始位置。

绿色(Mask Token):锚点 token 后面需要预测的位置,用 “<mask>” 标记。

核心逻辑:草稿模型的任务是 “以锚点 token 为起点,并行预测所有 mask 位置的 token”,和推理时 “以目标模型生成的 token 为锚点,并行生成草稿” 的场景完全对齐。

白色(Invisible Token):注意力掩码,作用是禁止不同草稿块之间互相 “偷看” 信息。

比如一个训练序列里有 3 个草稿块,白色掩码会让每个块只能关注自己内部的 token 和蓝色的 “目标思路”,不能跨块借鉴 —— 避免训练时作弊,确保草稿模型学会 “独立完成一个块的并行预测”。

image.png

实验结果与结论

实验1 —— no Reasoning SpeedUp

image.png

NVIDIA H200 GPU

DFlash训练数据:800K samples from NVIDIA Nemotron Post-Training Dataset V2 and CodeAlpaca. 没有直接使用原始数据集的Response,而是使用Target模型的Response作为训练集。

对于 DFlash 草稿模型,Decoder Layer 设置为 5(Qwen3 Coder 为 8),Block Size 为 16 的块大小(LLaMA 3.1 为 10)。目标隐藏特征是从目标模型的第二层到倒数第三层之间均匀选取的 5 个层中提取的。

实验中的Eagle3结果是基于AngelSlim的checkpoint。很遗憾的是没有找到AngelSlim训练Eagle3的数据量是多少。

推理后端:Transformers

<aside>

这里的加速比,不能直接和我们现在所做的实验进行比较。

我们实验的加速比的计算是基于完整的推理时间Prefilling + Decoding

DFlash加速比的计算只基于Decoding

</aside>

实验2 —— Reasoning SpeedUp