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

DFlash

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 和蓝色的 “目标思路”,不能跨块借鉴 —— 避免训练时作弊,确保草稿模型学会 “独立完成一个块的并行预测”。


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>