DFlash

DFlash (Chen et al., 2026) 是一个基于 块扩散(Block Diffusion) 的 speculative-decoding draft model：用一个 5 层轻量 Transformer 一次前向并行生成整个 token block（默认 16 tokens），实现常数级 drafting 延迟。核心创新是 KV 注入(KV Injection)——从目标模型的多层提取 hidden states，投影后注入到 draft model 每一层的 KV cache 中。在 Qwen3-8B 上达到约 4.86× 加速（SGLang 集成后 6.1×），比 EAGLE-3 快约 2.5×。

核心思想

传统 draft model（如 EAGLE 系列）以自回归方式逐 token 生成草稿，drafting 延迟与草稿长度 $\gamma$ 线性增长。DFlash 的关键洞察是：将 drafting 问题重新定义为去噪(Denoising)——给定一个初始 block（第一个 token 已知，其余为 mask token），一次前向预测整个 block 中所有位置的 token。因为 block 内所有位置同时计算，drafting 延迟是常数级的，不随 block size 增长。

与 EAGLE 的另一个关键区别在于目标模型特征的利用方式：

维度	EAGLE 系列	DFlash
Draft 方式	自回归（逐 token）	Block diffusion（并行整块）
目标特征注入位置	仅输入层（concat embedding）	每一层 KV cache
注意力类型	因果注意力(Causal)	非因果注意力(Non-Causal)
Drafting 延迟	$O(\gamma)$	$O(1)$
块内位置间信息	无（每个位置只看前文）	有（双向互相 attend）

Draft Head 架构

整体结构

DFlashDraftModel 继承自 Qwen3PreTrainedModel，由以下组件构成：

• layers：$N$ 层 Qwen3DFlashDecoderLayer（论文中 $N = 5$），每层包含 KV Injection attention + MLP
• fc：线性投影层，将多层目标特征拼接后降维：$(\text{num\_target\_layers} \times d_\text{hidden}) \to d_\text{hidden}$
• hidden_norm：RMSNorm，归一化投影后的目标特征
• rotary_emb：独立的 RoPE 模块，为 draft model 计算位置编码
• norm：最终的 RMSNorm，归一化输出 hidden states

Draft model 不包含自己的 embedding 层和 lm_head——它直接复用目标模型的 embed_tokens 和 lm_head。

输入与输出

输入：

参数	形状	含义
noise_embedding	$[B, \text{block\_size}, d_\text{hidden}]$	当前 block 的 token embedding（由目标模型 embed_tokens 生成）
target_hidden	$[B, L_\text{ctx}, n_\text{layers} \times d_\text{hidden}]$	从目标模型多层提取并拼接的 hidden states
position_ids	$[B, L_\text{ctx} + \text{block\_size}]$	上下文 + 当前 block 的位置编码索引

输出：

返回值	形状	含义
hidden states	$[B, \text{block\_size}, d_\text{hidden}]$	输出特征，经目标模型 lm_head 转换为 logits

注意输出形状与 noise_embedding 相同——draft model 只处理 block 内的 token，不输出上下文位置的 hidden states。上下文信息通过 KV Injection 参与注意力计算，但不出现在输出中。

forward 数据流

noise_embedding ─────────────────────── hidden_states
                                            │
target_hidden ── fc ── hidden_norm ──── target_hidden'
                                            │
                                    ┌───────┴───────┐
                                    │  position_ids  │
                                    │       │        │
                                    │  rotary_emb    │
                                    │       │        │
                                    │  cos, sin      │
                                    └───────┬───────┘
                                            │
                              ┌─────────────┼─────────────┐
                              │      Decoder Layer ×N      │
                              │  ┌─────────────────────┐   │
                              │  │  LN → KV Injection  │   │
                              │  │  Attention → Res.   │   │
                              │  │  LN → MLP → Res.    │   │
                              │  └─────────────────────┘   │
                              └─────────────┼─────────────┘
                                            │
                                    final RMSNorm
                                            │
                              target.lm_head (外部)
                                            │
                                        logits

KV Injection 注意力

Qwen3DFlashAttention 是 DFlash 最核心的模块。它将目标模型的特征以 KV 的形式注入到每一层注意力计算中。

计算过程

设当前 block 的 hidden states 为 $\mathbf{H} \in \mathbb{R}^{B \times L_q \times d}$，目标特征为 $\mathbf{C} \in \mathbb{R}^{B \times L_c \times d}$：

$\mathbf{Q} = \text{QNorm}(\mathbf{H} W_Q)$

$\mathbf{K} = \text{KNorm}\!\left(\text{concat}\!\left[\mathbf{C} W_K,\; \mathbf{H} W_K\right]\right)$

$\mathbf{V} = \text{concat}\!\left[\mathbf{C} W_V,\; \mathbf{H} W_V\right]$

$\text{Attn}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\!\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_h}}\right)\mathbf{V}$

关键设计点：

1. Q 仅来自 block 自身：$\mathbf{Q}$ 由 hidden_states 投影得到，长度为 $L_q$（block_size）
2. K/V 拼接上下文和自身：K 和 V 分别对 target_hidden（上下文）和 hidden_states（自身）做投影，然后在序列维度拼接，长度为 $L_c + L_q$
3. 非因果注意力：is_causal = False。block 内的位置可以双向 attend——位置 $i$ 可以看到位置 $j$（$j > i$），这让模型能利用块内所有位置的互信息来联合去噪
4. 位置编码对齐：由于 K 长度为 $L_c + L_q$ 而 Q 长度为 $L_q$，DFlash 修改了标准 apply_rotary_pos_emb——K 使用完整的位置 cos/sin（前 $L_c$ 个对应上下文位置，后 $L_q$ 个对应 block 位置），Q 只取最后 $L_q$ 个

# 修改后的 RoPE：Q 和 K 使用不同长度的位置编码
q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :])  # 只取后 q_len 个
k_embed = (k * cos) + (rotate_half(k) * sin)                                     # 使用全部

与标准 Cross-Attention 的区别

KV Injection 不是简单的 cross-attention。标准 cross-attention 的 KV 完全来自外部源，而 KV Injection 的 KV 是上下文特征和自身特征的拼接。这意味着 block 内的每个位置既能看到目标模型提供的上下文信息（通过 $\mathbf{C} W_K, \mathbf{C} W_V$），也能看到 block 内其他位置的信息（通过 $\mathbf{H} W_K, \mathbf{H} W_V$）。后者对于块内联合去噪至关重要。

详细的代码与公式逐行对照见 KV Injection 完整数据流 Walkthrough，其中用一个 3-token prompt + block_size=4 的例子走通了从 hidden states 提取、RoPE 对齐到注意力矩阵计算的每一步。

Hidden States 提取与使用

从目标模型提取

DFlash 需要从目标模型的多个中间层提取 hidden states，而不只是使用最终输出。提取逻辑涉及两个工具函数：

1. build_target_layer_ids(num_target_layers, num_hidden_layers)：在目标模型的 num_hidden_layers 层中均匀采样 num_target_layers 个层的索引。例如，对 32 层的 Qwen3-8B 取 5 层，大约得到第 3、9、16、22、28 层（均匀间隔）

2. extract_context_feature(output.hidden_states, target_layer_ids)：从目标模型的 output_hidden_states（一个 tuple，包含每层的 $[B, L, d_\text{hidden}]$ 张量）中取出指定层的 hidden states，在特征维度上拼接，得到 $[B, L, n_\text{layers} \times d_\text{hidden}]$

目标模型必须以 output_hidden_states=True 运行才能返回中间层特征。

投影与归一化

提取的多层特征先经过线性投影降维，再归一化：

$\mathbf{C} = \text{RMSNorm}\!\left(W_\text{fc}\;\text{concat}\!\left[\mathbf{h}^{(l_1)}, \mathbf{h}^{(l_2)}, \dots, \mathbf{h}^{(l_k)}\right]\right)$

其中 $W_\text{fc} \in \mathbb{R}^{(k \cdot d) \times d}$ 是无 bias 的线性层。投影后的 $\mathbf{C}$ 被送入 draft model 的每一层 KV Injection attention，作为所有层共享的上下文信号。

在推理循环中的使用

在 spec_generate 中，hidden states 提取发生在两个时机：

Prefill 阶段：

output = target(input_ids, ..., output_hidden_states=True)
target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids)
# target_hidden: [B, num_input_tokens, num_target_layers * d_hidden]

首次提取覆盖整个 prompt 的所有位置。这为第一个 decode block 提供完整上下文。

每轮验证后：

output = target(block_output_ids, ..., output_hidden_states=True)
target_hidden = extract_context_feature(
    output.hidden_states, self.target_layer_ids
)[:, :acceptance_length + 1, :]

每次 target model 验证一个 block 后，从验证输出中提取 hidden states，但只保留被接受位置的特征（[:acceptance_length + 1]）。这些新特征在下一轮 draft 中作为上下文注入。

Draft model 的 KV cache 通过 past_key_values_draft.crop(start) 保持与已接受序列同步——每轮丢弃当前 block 的 KV，只保留上下文部分。这样 draft model 的上下文逐轮增长：第一轮看到 prompt 特征，第二轮追加第一轮被接受的 token 的特征，以此类推。

推理流程

spec_generate 实现了完整的投机解码主循环：

1. 初始化

output_ids = torch.full((1, max_length + block_size), mask_token_id, ...)

预分配一个全 mask 的输出缓冲区。mask_token_id 标记尚未生成的位置。

2. Prefill

对 prompt 做一次目标模型前向：获取首个生成 token + 全量上下文特征。

3. Decode 循环

每轮迭代：

1. 准备 block：从 output_ids 取出 [start, start + block_size) 的 token 作为当前 block（首位为上一轮最后接受的 token，其余为 mask 或上一轮残留）
2. Draft：将 block token 通过目标模型 embed_tokens 得到 noise embedding，与目标特征一起送入 draft model，取后 block_size - 1 个位置的 logits 采样得到草稿 token
3. Verify：将填入草稿 token 的 block 送入目标模型，获取 posterior 分布
4. Accept/Reject：逐位比较草稿 token 与 posterior 采样结果，通过 cumprod 找到最长匹配前缀作为接受长度。第一个不匹配位置用 posterior 的采样结果替换
5. 更新状态：裁剪两个 KV cache（target 和 draft），提取新的上下文特征，推进 start 位置

为什么取 [:, -block_size+1:, :]

Draft model 输出 block_size 个位置的 hidden states，但第 0 个位置对应的 token 已经是已知的（上一轮最后接受的 token），不需要预测。因此只取位置 1 到 block_size-1 的输出，经 lm_head 得到 block_size - 1 个 logits，分别预测 block 中第 1 到第 block_size-1 个位置应该是什么 token。

相关页面

• speculative-decoding — 投机解码的原理、验证算法与 draft model 设计流派
• kv-cache — KV cache 的算法作用与工程实现
• eagle-mtp-online-training — EAGLE MTP 在 RL 场景下的在线训练