DFlash KV Injection Walkthrough

本页是 dflash 的补充材料，通过一个具体数值示例，逐行对照代码与数学公式，完整走通一轮 draft 中从 hidden states 提取到 KV Injection 注意力计算的全过程。

示例设定

参数	值	说明
Target model 层数	12
$d_\text{hidden}$	512
$n_h$ (注意力头数)	8	$d_h = d / n_h = 64$
Draft model 层数 $N$	2
提取层数 $k$	3	从 target 的第 3/7/11 层提取
block_size	4
Prompt	"The cat sat"	3 tokens，位置 0/1/2

为简化，示例中省略 batch 维度。多头注意力的公式按单头写，实际实现中对 $n_h$ 个头独立执行后拼接。

Step 1：Prefill — 从 target model 提取 hidden states

Target model 以 output_hidden_states=True 处理 prompt：

output = target(
    input_ids,                              # ["The", "cat", "sat"]
    position_ids=position_ids[:, :3],       # [0, 1, 2]
    past_key_values=past_key_values_target,
    use_cache=True,
    logits_to_keep=1,
    output_hidden_states=True,              # ← 关键：返回所有层的 hidden states
)

output.hidden_states 是一个长度为 13 的 tuple（embedding 层 + 12 个 Transformer 层），每个元素形状为 $[3, 512]$。extract_context_feature 选取第 3/7/11 层，在特征维拼接：

                    位置 0 ("The")  位置 1 ("cat")  位置 2 ("sat")
Target 第 3 层 :      h³₀             h³₁             h³₂            ← 各为 512 维
Target 第 7 层 :      h⁷₀             h⁷₁             h⁷₂
Target 第 11 层:      h¹¹₀            h¹¹₁            h¹¹₂

$\text{target\_hidden\_raw}[i] = \left[\mathbf{h}^{(3)}_i ;\; \mathbf{h}^{(7)}_i ;\; \mathbf{h}^{(11)}_i\right] \in \mathbb{R}^{1536}$

target_hidden = extract_context_feature(output.hidden_states, self.target_layer_ids)
# target_hidden: [3, 1536]   (3 个位置，每个位置 512 × 3 = 1536 维)

Step 2：投影与归一化

在 DFlashDraftModel.forward 的入口处，多层拼接特征被投影回模型维度：

target_hidden = self.hidden_norm(self.fc(target_hidden))

对应数学：

$\mathbf{C} = \text{RMSNorm}\!\left(W_\text{fc} \cdot \text{target\_hidden\_raw}\right)$

其中 $W_\text{fc} \in \mathbb{R}^{1536 \times 512}$（无 bias），将三层信息压缩为：

$\mathbf{C} \in \mathbb{R}^{3 \times 512}$

这个 $\mathbf{C}$ 在 draft model 的所有层中作为同一份上下文信号使用。

Step 3：准备 block

Prefill 同时采样了首个新 token "on"（位置 3）。当前 block = ["on", MASK, MASK, MASK]（位置 3/4/5/6）：

block_output_ids = output_ids[:, start : start + block_size].clone()
# block_output_ids = ["on", MASK, MASK, MASK]

noise_embedding = target.model.embed_tokens(block_output_ids)
# noise_embedding: [4, 512]   即 H ∈ R^{4×512}

Step 4：KV Injection 注意力——代码与公式逐行对照

以下逐行走通 Qwen3DFlashAttention.forward。设 $L_c = 3$（上下文），$L_q = 4$（block），$d_h = 64$（每头维度）。

4.1 Q 投影

q = self.q_proj(hidden_states)                    # [4, 512] → [4, 512]
q = q.view(bsz, q_len, -1, self.head_dim)         # [4, 512] → [4, 8, 64]
q = self.q_norm(q).transpose(1, 2)                 # QNorm → [8, 4, 64]

$\mathbf{Q} = \text{QNorm}\!\left(\mathbf{H} W_Q\right) \in \mathbb{R}^{n_h \times L_q \times d_h} = \mathbb{R}^{8 \times 4 \times 64}$

Q 仅来自 block 自身的 hidden_states（即 noise embedding 经过前面层处理后的结果），共 4 个 query 向量，对应位置 3/4/5/6。

4.2 K/V 投影与拼接

这是 KV Injection 的核心——K 和 V 分别对上下文和 block 自身做投影，然后在序列维拼接：

k_ctx   = self.k_proj(target_hidden)     # C·W_K: [3, 512] → [3, 512]
k_noise = self.k_proj(hidden_states)     # H·W_K: [4, 512] → [4, 512]
v_ctx   = self.v_proj(target_hidden)     # C·W_V: [3, 512] → [3, 512]
v_noise = self.v_proj(hidden_states)     # H·W_V: [4, 512] → [4, 512]

k = torch.cat([k_ctx, k_noise], dim=1)  # [7, 512] → view → [7, 8, 64]
v = torch.cat([v_ctx, v_noise], dim=1)  # [7, 512] → view → [7, 8, 64]
k = self.k_norm(k).transpose(1, 2)      # KNorm → [8, 7, 64]
v = v.transpose(1, 2)                   # [8, 7, 64]

$\mathbf{K} = \text{KNorm}\!\left(\text{concat}\!\left[\mathbf{C} W_K,\; \mathbf{H} W_K\right]\right) \in \mathbb{R}^{n_h \times (L_c + L_q) \times d_h} = \mathbb{R}^{8 \times 7 \times 64}$

$\mathbf{V} = \text{concat}\!\left[\mathbf{C} W_V,\; \mathbf{H} W_V\right] \in \mathbb{R}^{8 \times 7 \times 64}$

7 个 KV 向量的结构：

K/V 索引	位置	来源	含义
0	pos 0	$\mathbf{C} W_K$	target 对 "The" 的多层融合特征
1	pos 1	$\mathbf{C} W_K$	target 对 "cat" 的多层融合特征
2	pos 2	$\mathbf{C} W_K$	target 对 "sat" 的多层融合特征
3	pos 3	$\mathbf{H} W_K$	block 中 "on" 的 embedding
4	pos 4	$\mathbf{H} W_K$	block 中 MASK 的 embedding
5	pos 5	$\mathbf{H} W_K$	block 中 MASK 的 embedding
6	pos 6	$\mathbf{H} W_K$	block 中 MASK 的 embedding

注意 $\mathbf{C} W_K$ 和 $\mathbf{H} W_K$ 使用的是同一个 k_proj 权重矩阵——上下文和 block 自身的 key 在同一个投影空间中。V 同理。

4.3 RoPE 位置编码对齐

Q 有 4 个向量，K 有 7 个向量——长度不同。标准 RoPE 要求 Q 和 K 长度一致，DFlash 修改了 apply_rotary_pos_emb 来处理这个差异：

cos, sin = position_embeddings
# rotary_emb 对 position_ids = [0,1,2,3,4,5,6] 计算
# cos, sin: [7, 64]

q, k = apply_rotary_pos_emb(q, k, cos, sin)

apply_rotary_pos_emb 的实现：

def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)   # [1, 7, 64]
    sin = sin.unsqueeze(unsqueeze_dim)   # [1, 7, 64]
    q_len = q.size(-2)                   # = 4

    # Q 只取最后 q_len 个位置的 cos/sin
    q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :])
    #                    ↑ cos[..., 3:7, :] 即位置 3,4,5,6

    # K 使用全部 7 个位置的 cos/sin
    k_embed = (k * cos) + (rotate_half(k) * sin)
    #              ↑ cos[..., 0:7, :] 即位置 0,1,2,3,4,5,6

    return q_embed, k_embed

RoPE 的作用是让内积 $\mathbf{q}_i^\top \mathbf{k}_j$ 只依赖于相对位置 $i - j$。设 $R_m$ 为位置 $m$ 的旋转矩阵：

$\mathbf{q}'_i = R_i \mathbf{q}_i, \quad \mathbf{k}'_j = R_j \mathbf{k}_j$

$\mathbf{q}'^{\top}_i \mathbf{k}'_j = \mathbf{q}_i^\top R_i^\top R_j \mathbf{k}_j = \mathbf{q}_i^\top R_{j-i} \mathbf{k}_j$

在本例中对齐情况：

Q: q₃  q₄  q₅  q₆         → R₃, R₄, R₅, R₆
K: k₀  k₁  k₂  k₃  k₄  k₅  k₆  → R₀, R₁, R₂, R₃, R₄, R₅, R₆

例如 $q_5^\top k_2$：旋转角差 $= R_{2-5} = R_{-3}$，正确编码了「位置 5 的 query 看位置 2 的 key，距离为 3」。

如果错误地让 Q 使用 $[0,1,2,3]$ 而非 $[3,4,5,6]$，则 $q_5$ 会被编码为位置 2，$q_5^\top k_2$ 的旋转角差变为 $R_0$——模型会误认为它们在同一位置。

4.4 注意力计算

attn_output, attn_weights = attn_fn(
    self, q, k, v, attention_mask,
    dropout=0.0 if not self.training else self.attention_dropout,
    scaling=self.scaling,     # = d_h^{-0.5} = 64^{-0.5} = 0.125
    sliding_window=self.sliding_window,
    **kwargs,
)

由于 is_causal = False，注意力矩阵没有因果 mask，是一个 $4 \times 7$ 的全连接矩阵：

$A_{ij} = \frac{\exp\!\left(\mathbf{q}'^\top_i \mathbf{k}'_j / \sqrt{d_h}\right)}{\sum_{m=0}^{6} \exp\!\left(\mathbf{q}'^\top_i \mathbf{k}'_m / \sqrt{d_h}\right)}, \quad i \in \{3,4,5,6\},\; j \in \{0,\dots,6\}$

$\text{out}_i = \sum_{j=0}^{6} A_{ij} \mathbf{v}_j$

展开位置 5（MASK）的计算：

$\text{out}_5 = \underbrace{A_{5,0}\mathbf{v}_0 + A_{5,1}\mathbf{v}_1 + A_{5,2}\mathbf{v}_2}_{\text{来自 target 上下文（KV Injection）}} + \underbrace{A_{5,3}\mathbf{v}_3 + A_{5,4}\mathbf{v}_4 + A_{5,5}\mathbf{v}_5 + A_{5,6}\mathbf{v}_6}_{\text{来自 block 自身（双向自注意力）}}$

注意力矩阵的可视化（每个 ● 代表一个 $A_{ij}$ 权重）：

          K:  pos0    pos1    pos2   ┃  pos3    pos4    pos5    pos6
              "The"   "cat"   "sat"  ┃  "on"    MASK    MASK    MASK
              ← target 上下文特征 →  ┃  ←─── block 自身 embedding ──→
                                     ┃
Q: pos3 "on"    ●       ●       ●   ┃    ●       ●       ●       ●
   pos4 MASK    ●       ●       ●   ┃    ●       ●       ●       ●
   pos5 MASK    ●       ●       ●   ┃    ●       ●       ●       ●
   pos6 MASK    ●       ●       ●   ┃    ●       ●       ●       ●
                                     ┃
                             KV Injection 边界
                     左侧 V 来自 target 特征
                     右侧 V 来自 block embedding

对比：如果使用因果注意力（EAGLE 的方式）：

          K:  pos3    pos4    pos5    pos6
              "on"     ?       ?       ?

Q: pos3       ●       ×       ×       ×
   pos4       ●       ●       ×       ×
   pos5       ●       ●       ●       ×
   pos6       ●       ●       ●       ●

× = 因果 mask 遮挡

EAGLE 中 pos5 只能看到 pos3 和 pos4 的信息。KV Injection 中 pos5 能同时看到：

• target model 对 "The cat sat" 的深层语义理解（pos0-2 的 $\mathbf{v}$）
• 已知 token "on" 的信息（pos3 的 $\mathbf{v}$）
• 其他 MASK 位置的互信息（pos4, pos6 的 $\mathbf{v}$，如果前一层已部分去噪）

4.5 输出投影

attn_output = attn_output.reshape(bsz, q_len, -1)  # [4, 8, 64] → [4, 512]
attn_output = self.o_proj(attn_output)               # [4, 512] → [4, 512]

$\text{Attn\_out} = \text{concat}[\text{head}_1, \dots, \text{head}_{n_h}] \cdot W_O \in \mathbb{R}^{L_q \times d} = \mathbb{R}^{4 \times 512}$

注意输出形状是 $[4, 512]$——与输入 hidden_states 相同，只有 block 的 4 个位置。上下文的 3 个位置只参与了 KV 计算，不产生输出。

Step 5：输出与 logits

经过 2 层 Decoder Layer + final RMSNorm 后，draft model 输出 $[4, 512]$。

取后 3 个位置（[:, -block_size+1:, :] = 位置 4/5/6 的输出），经 target.lm_head 得到 3 组 logits：

draft_logits = target.lm_head(self(...)[:, -block_size+1:, :])
# self(...): [4, 512]
# [:, -3:, :]: [3, 512]    取位置 4, 5, 6
# lm_head:    [3, vocab_size]

block_output_ids[:, 1:] = sample(draft_logits)

位置 3（"on"）的输出被丢弃——它已是已知 token，不需要预测。

block 更新前: ["on",  MASK,  MASK,  MASK]   位置 3, 4, 5, 6
                      ↓      ↓      ↓
draft 采样:          "the"  "mat"   "."
block 更新后: ["on", "the", "mat",  "."]

Step 6：验证与接受

Target model 对整个 block 做一次前向：

output = target(
    block_output_ids,                        # ["on", "the", "mat", "."]
    position_ids=block_position_ids,         # [3, 4, 5, 6]
    past_key_values=past_key_values_target,
    use_cache=True,
    output_hidden_states=True,
)
posterior = sample(output.logits, temperature)

Target 的自回归 logits 给出每个位置下一个 token 的预测：

pos 3 "on"  → target 预测 pos 4 = "the"  ← 与 draft pos4 "the" 匹配 ✓
pos 4 "the" → target 预测 pos 5 = "mat"  ← 与 draft pos5 "mat" 匹配 ✓
pos 5 "mat" → target 预测 pos 6 = "in"   ← 与 draft pos6 "."  不匹配 ✗
pos 6 "."   → (不再检查)

验证代码：

acceptance_length = (block_output_ids[:, 1:] == posterior[:, :-1]).cumprod(dim=1).sum(dim=1)[0].item()

逐步展开：

$\text{block\_output\_ids}[:, 1:] = [\text{"the"},\; \text{"mat"},\; \text{"."}]$

$\text{posterior}[:, :-1] = [\text{"the"},\; \text{"mat"},\; \text{"in"}]$

$\text{match} = [1, 1, 0] \xrightarrow{\text{cumprod}} [1, 1, 0] \xrightarrow{\text{sum}} 2$

acceptance_length = 2。最终序列：

..., "sat", "on", "the", "mat", "in", ...
                                  ↑
                        posterior 修正（替换 draft 的 "."）

Step 7：更新状态，进入下一轮

start += acceptance_length + 1                      # start = 3 + 3 = 6 → 下一个 block 从 pos 6 开始
                                                     # 实际 start 前进到 pos 7（"in" 的下一位），
                                                     # 因为 "in" 被写入 pos 6

past_key_values_target.crop(start)                   # target KV cache 裁剪到 pos 0..6
target_hidden = extract_context_feature(
    output.hidden_states, self.target_layer_ids
)[:, :acceptance_length + 1, :]                      # 只取 pos 3,4,5 的特征（被接受的 3 个位置）

Draft model 的 KV cache 在本轮 draft 结束时已经 crop 到 start（只保留上下文 KV）。下一轮进入时，新的 target_hidden（pos 3/4/5 的特征）会作为增量上下文注入。Draft model 的 KV cache 逐轮增长：

第 1 轮: KV cache 包含 pos 0,1,2 的 KV（prompt 的 target 特征）
第 2 轮: KV cache 包含 pos 0,1,2,3,4,5 的 KV（+ 本轮接受的 3 个位置）
第 3 轮: KV cache 包含 pos 0,...,N 的 KV（继续增长）