03 · DSpark 建模：Markov 与 Confidence

本篇是 DSpark 实现的核心拆解，对应论文 Section 3.1（Semi-Autoregressive Generation）、3.2.1（Confidence Head）、3.3（Training）。DSpark 是 DeepSpec 的主角算法，本篇从 forward 13 步流程、anchor 采样、block attention mask、三种 Markov head、confidence head 输入构造、loss 三项加权 6 个维度逐一拆解。配套 04 Eagle3 对照自回归范式。

---

总览段（总）

DSpark 的设计目标是用一次并行 forward 产出 block_size 个 draft token，同时通过极轻的串行 Markov head 注入前缀依赖，再用 confidence head 预测每个位置的存活概率。整个 forward 在 Qwen3DSparkModel.forward 中实现，13 个步骤从 anchor 采样到 DSparkForwardOutput 返回。

图说明： DSpark forward 是一个 13 步流水线，对应 [qwen3/modeling.py:389-526](file:///workspace/deepspec/modeling/dspark/qwen3/modeling.py#L389-526)。输入 4 个张量（input_ids / target_hidden_states / loss_mask / target_last_hidden_states），输出 DSparkForwardOutput 含 6 个张量。关键创新在第 4 步（block attention mask，让 context 与 draft block 在同一 forward 中互不污染）、第 12 步（Markov head 注入前缀依赖）、第 13 步（confidence head 预测接受率）。

关键文件清单：

文件	角色
[deepspec/modeling/dspark/common.py](file:///workspace/deepspec/modeling/dspark/common.py)	公共组件：anchor 采样 / block mask / noise embed / eval_mask / AcceptRatePredictor
[deepspec/modeling/dspark/markov_head.py](file:///workspace/deepspec/modeling/dspark/markov_head.py)	三种 Markov head：Vanilla / Gated / RNN
[deepspec/modeling/dspark/loss.py](file:///workspace/deepspec/modeling/dspark/loss.py)	CE / L1 / Confidence BCE 三项 loss
[deepspec/modeling/dspark/qwen3/modeling.py](file:///workspace/deepspec/modeling/dspark/qwen3/modeling.py)	Qwen3 后端 DSpark 实现
[deepspec/modeling/dspark/qwen3/config.py](file:///workspace/deepspec/modeling/dspark/qwen3/config.py)	Qwen3 draft config 构造
[deepspec/modeling/dspark/gemma4/modeling.py](file:///workspace/deepspec/modeling/dspark/gemma4/modeling.py)	Gemma4 后端（与 Qwen3 对称）

---

分述段（分）

3.1 DSparkForwardOutput 数据结构

[DSparkForwardOutput](file:///workspace/deepspec/modeling/dspark/common.py) （[common.py:12-40](file:///workspace/deepspec/modeling/dspark/common.py#L12-40)）：

图说明： 形状记号 B=batch、A=num_anchors、γ=block_size、V=vocab_size、H=hidden_size。eval_mask 用 cumprod(dim=-1) 强制前缀连续——一旦某位置无效，其后所有位置自动归零（[common.py:172-188](file:///workspace/deepspec/modeling/dspark/common.py#L172-188)）。block_keep_mask 标记 anchor 是否真实采样（不足 num_anchors 时填充 dummy）。aligned_target_logits 仅在训练时提供（用于 L1 loss 与 accept rate 监督），评测时为 None。

3.2 Anchor 采样：训练数据的切分方式

[sample_anchor_positions](file:///workspace/deepspec/modeling/dspark/common.py) （[common.py:123-169](file:///workspace/deepspec/modeling/dspark/common.py#L123-169)）：

• 候选 mask：build_anchor_candidate_mask（[common.py:109-120](file:///workspace/deepspec/modeling/dspark/common.py#L109-120)）——anchor 候选位置 = 当前位置和下一位置都在 loss_mask 内的位置（保证 anchor 后还能取 γ 个有效 label）。
• 采样：对每个样本最多采样 num_anchors=512 个 anchor，采用随机值排序后取前 N 个并按位置升序排列。
• 不足时填充：dummy anchor 用 block_keep_mask=0 屏蔽，不参与 loss。

这与论文 Section 3.3 描述完全一致：“we randomly sample multiple anchor positions from each target sequence to form γ-token blocks as training data”。

3.3 Noise Embedding：anchor + γ-1 mask tokens

[create_noise_embed](file:///workspace/deepspec/modeling/dspark/common.py) （[common.py:264-294](file:///workspace/deepspec/modeling/dspark/common.py#L264-294)）：

图说明： 每个 block 的输入是 1 个 anchor token（位置 p 的真实 token）+ γ-1 个 mask token（mask_token_id=151669 是 Qwen3 的 <|mask|>）。这是论文 Section 3.1 “Parallel stage” 的代码落地——“γ input tokens (anchor + γ-1 masks) yield γ draft logits”。注意 DSpark 把 anchor 本身也作为第一个预测位置，相比 DFlash 原版（anchor 不预测）少一次 forward 计算。

3.4 Block Attention Mask：并行 forward 的隔离墙

[create_dspark_attention_mask](file:///workspace/deepspec/modeling/dspark/common.py) （[common.py:78-106](file:///workspace/deepspec/modeling/dspark/common.py#L78-106)）用 torch.nn.attention.flex_attention.create_block_mask 构造一个混合 mask：

图说明： 这是 DSpark 的核心创新之一。整个序列是 context（target hidden，长度 seq_len）拼接所有 draft block（长度 num_blocks × γ），用单一 forward 同时处理。规则保证：① context 内标准 causal；② 每个 draft block 只能看自己 anchor 之前的 context + 同 block 内的 KV，绝不漏看其他 block 的 draft token；③ 无效 block（block_keep_mask=0）整块屏蔽。这避免了 padding 浪费，让 num_anchors 个 block 一次 forward 出全部 logits。代码定义在 dspark_mask_mod（[common.py:86-96](file:///workspace/deepspec/modeling/dspark/common.py#L86-96)）。

3.5 Custom Attention：context K/V 与 draft K/V 拼接

Qwen3DSparkAttention（[qwen3/modeling.py:44-152](file:///workspace/deepspec/modeling/dspark/qwen3/modeling.py#L44-152)）的关键差异：forward 接受两路输入——target_hidden_states（context）与 hidden_states（noise embedding，draft）。K/V 投影同时作用于两路再拼接：

k = cat([k_proj(target), k_proj(noise)], dim=1)   # qwen3/modeling.py:108-113
v = cat([v_proj(target), v_proj(noise)], dim=1)

这实现了"cross-attention to target + self-attention over draft"的混合模式。Q 来自 draft noise embedding，所以 query 是"基于 anchor + mask tokens 的 draft 表征"，K/V 是"context + draft"两部分拼接。

3.6 build_eval_mask：cumprod 强制前缀连续

[build_eval_mask](file:///workspace/deepspec/modeling/dspark/common.py) （[common.py:172-188](file:///workspace/deepspec/modeling/dspark/common.py#L172-188)）：

图说明： 三个条件 AND 后做 cumprod(dim=-1)——一旦某位置为 0，其后所有位置自动归零。这与推测解码"接受最长正确前缀"的语义完全对齐：训练时若位置 k 应被屏蔽（如超出 loss_mask 范围），位置 k+1, k+2, … 全部归零，loss 只算连续有效前缀。

3.7 三种 Markov head

[build_markov_head](file:///workspace/deepspec/modeling/dspark/markov_head.py) （[markov_head.py:287-311](file:///workspace/deepspec/modeling/dspark/markov_head.py#L287-311)）工厂函数，根据 markov_head_type 配置返回三种之一。markov_rank=0 时返回 None（DFlash 走这条路径）。

图说明： 三种 head 对应论文 Section 3.1 的三个公式：

VanillaMarkov（[markov_head.py:8-90](file:///workspace/deepspec/modeling/dspark/markov_head.py#L8-90)，论文 Eq. 5）：参数 W1 ∈ R^{V×r}、W2 ∈ R^{r×V}，bias = W2(W1[prev_token])。低秩分解 r=256 让参数量从 $V^2$（Qwen3 词表 ~150K，约 22.5B 参数）降到 $2rV$（~77M）。apply_block_logits 训练时 teacher-forced，每个位置的 bias 由对应 prev token 决定；sample_block_tokens 推理时按序采样，prev_token_ids 迭代更新。

GatedMarkovHead（[markov_head.py:93-122](file:///workspace/deepspec/modeling/dspark/markov_head.py#L93-122)）：在 vanilla 基础上加 gate，让 backbone hidden state 调制 Markov 偏置：

$$\text{gate}=\sigma (W_g[h;W_1[x_{k-1}]])，B=W_2(\text{gate}\odot W_1[x_{k-1}])$$

RNNHead（[markov_head.py:125-284](file:///workspace/deepspec/modeling/dspark/markov_head.py#L125-284)，论文 Eq. 6）：GRU-like 循环状态，拼接 $z_k=[s_{k-1};W_1[x_{k-1}];h_k]\in R^{2r+d}$，单层 gated update：

$$s_k=\sigma (W_g{z}_k)\odot s_{k-1}+(1-\sigma (W_g{z}_k))\odot \tanh (W_c{z}_k),B_k=W_2^{\top }\tanh (W_o{z}_k)$$

apply_block_logits 训练时 teacher-forced 但状态在 block 内累积；sample_block_tokens 推理时状态在采样步之间传递。论文 Section 4.3.2 实测 RNN head 相比 Markov head 只在长 proposal length（γ=12,16）有边际增益，默认用 Markov。

3.8 Confidence head：预测前缀接受率

[AcceptRatePredictor](file:///workspace/deepspec/modeling/dspark/common.py) （[common.py:43-49](file:///workspace/deepspec/modeling/dspark/common.py#L43-49)）：nn.Linear(input_dim, 1) + squeeze，输入维度由 confidence_head_with_markov 决定——若为 True 则 input_dim = hidden_size + markov_rank（拼 backbone hidden 与 markov embedding），否则 input_dim = hidden_size。

图说明： 对应论文 Eq. 7。confidence head 输出每个 draft 位置的"条件接受率"——给定前缀全被接受，本位置被接受的概率。监督目标是 accept_rate_3d.detach()（解析接受率，由 draft 与 target 分布的 total variation 距离算出）。代码中 [qwen3/modeling.py:505-517](file:///workspace/deepspec/modeling/dspark/qwen3/modeling.py#L505-517) 的输入构造体现了"with_markov"差异：若开启则把 markov head 的 get_prev_embeddings(prev_token_ids) 与 backbone hidden 拼接，让 confidence 看到前一个 draft token 的信息。

3.9 Loss 三项加权

[compute_dspark_loss](file:///workspace/deepspec/modeling/dspark/loss.py) （[loss.py:255](file:///workspace/deepspec/modeling/dspark/loss.py#L255)）汇总三项：

图说明： 对应论文 Eq. 9-12。三项 loss 都用位置权重 w_k = exp(-(k-1)/γ) 加权（loss_decay_gamma=4.0，[loss.py:25-37](file:///workspace/deepspec/modeling/dspark/loss.py#L25-37) 的 _build_loss_weight_mask），让靠前位置权重更大。关键细节：backward_loss 在 [loss.py:252](file:///workspace/deepspec/modeling/dspark/loss.py#L252) 乘以 world_size——因为梯度只在 local batch 上累积，需要按 world_size 缩放以匹配全局平均（_all_reduce_loss_denominators 在 [loss.py:11-22](file:///workspace/deepspec/modeling/dspark/loss.py#L11-22) 做 all_reduce SUM 把各 rank 的分母聚合）。

CE 损失（[loss.py:109-114](file:///workspace/deepspec/modeling/dspark/loss.py#L109-114)）：标准 F.cross_entropy，每个 draft 位置加权求和后除以全局分母。

L1 / TV 损失（[loss.py:73-87](file:///workspace/deepspec/modeling/dspark/loss.py#L73-87)）：l1_dist = |softmax(draft) - softmax(target)|.sum(-1)，对应论文 Eq. 10。这是 L1 距离 = 2 × total variation distance，与接受率直接对偶（论文 Section 3.3：“minimizing Ltv directly maximizes the expected acceptance rate”）。

Accept rate 3D（[loss.py:60-70](file:///workspace/deepspec/modeling/dspark/loss.py#L60-70)）：accept_rate = 1 - 0.5 * |draft_probs - target_probs|.sum(-1)，clamp 到 [0,1]。对应论文 Eq. 8。

Confidence BCE（[loss.py:152-181](file:///workspace/deepspec/modeling/dspark/loss.py#L152-181)）：F.binary_cross_entropy_with_logits(confidence_pred, accept_rate.detach()) 加权。同时记录 confidence_abs_error、confidence_bias、confidence_cumprod_bias 用于校准评估。

3.10 τ 指标：期望接受 draft 数

[_compute_local_probabilistic_stats](file:///workspace/deepspec/modeling/dspark/loss.py) （[loss.py:40-57](file:///workspace/deepspec/modeling/dspark/loss.py#L40-57)）：

expected_draft_accepted = (accept_rate * eval_mask).cumprod(dim=-1).sum(dim=-1)
tau_prob_per_block = expected_draft_accepted + 1   # +1 是 bonus token

这是论文 $\tau \approx 1+{\sum }_k{\prod }_{i\le k}{c}_i$ 的代码实现。训练时记录 tau_greedy 与 tau_probabilistic 两个指标——前者用 argmax token 算接受率，后者用概率分布算。这是训练日志的核心指标，可直接反映"训练了多少步后推测解码会变快"。

3.11 Aligned target logits 对齐策略

训练时若提供 target_last_hidden_states，需要把 target 的 hidden 对齐到 draft 的预测位置（[qwen3/modeling.py:448-466](file:///workspace/deepspec/modeling/dspark/qwen3/modeling.py#L448-466)）：

• target_pred_indices = (safe_label_indices - 1).clamp(min=0)：因为 draft 在位置 p 预测 token p+1，要看的是 target 在位置 p 的 hidden
• 用 torch.gather 抽取对齐后的 hidden
• 过 compute_logits 得到 aligned_target_logits，用于 L1 loss 与 accept rate 监督

3.12 Gemma4 后端的差异

Gemma4DSparkModel（[gemma4/modeling.py](file:///workspace/deepspec/modeling/dspark/gemma4/modeling.py)）与 Qwen3 版本对称，差异主要在：

• Gemma4DSparkAttention 使用 global_head_dim、num_global_key_value_heads、attention_k_eq_v（k=v 复用）、v_norm、scaling=1.0
• 使用 flex_attention 直接调用（而非 ALL_ATTENTION_FUNCTIONS）
• compute_logits 应用 final_logit_softcapping（tanh(logits/softcap)*softcap，[gemma4/modeling.py:340-349](file:///workspace/deepspec/modeling/dspark/gemma4/modeling.py#L340-349)）
• 嵌入层是 Gemma4TextScaledWordEmbedding，带 embed_scale = sqrt(hidden_size)
• Decoder layer 有额外的 pre_feedforward_layernorm / post_feedforward_layernorm 和 layer_scalar
• config 路径多一层 text_config（[gemma4/config.py:9-19](file:///workspace/deepspec/modeling/dspark/gemma4/config.py#L9-19)）

forward 主体逻辑（[gemma4/modeling.py:451-598](file:///workspace/deepspec/modeling/dspark/gemma4/modeling.py#L451-598)）与 Qwen3 完全一致——这印证了 DSpark 算法与 target 模型家族无关。

3.13 DSpark vs DFlash：配置差异即算法差异

回顾 01 架构 的发现，DFlash 在本仓库是 DSpark 的退化配置：

字段	DSpark	DFlash	含义
markov_rank	256	0	关闭串行 Markov head
markov_head_type	‘vanilla’	-	build_markov_head 返回 None
confidence_head_alpha	1.0	0.0	关闭 confidence 监督
confidence_head_with_markov	True	-	-
ce_loss_alpha	0.1	1.0	纯 CE 训练
l1_loss_alpha	0.9	0.0	关闭 TV 对齐

这印证了论文 Section 4.3.1 的诊断：DFlash suffix decay 来自并行独立性，DSpark 通过 Markov head 注入前缀依赖解决之；通过 confidence head 让验证变智能。

---

小结段（总）

DSpark 的工程实现可以浓缩为三句话：一次 forward 出 γ 个 draft token（并行主干）+ 极轻 Markov head 注入前缀依赖（解决 suffix decay）+ confidence head 预测接受率（让验证智能截断）。13 步 forward 流程把这三件事用 flex_attention block mask 一次性完成，无 padding、无串行 forward、无 target 模型在线推理。

设计要点回顾：

1. anchor 采样让 batch 内可同时监督多个 block：num_anchors=512，单 batch 最多 512×7=3584 个监督位置。
2. block mask 是核心创新：让 context 与 draft block 共用一次 forward，但保证 draft block 间互不污染。
3. Markov head 低秩分解：r=256 把 $V^2$ 参数降到 $2rV$，让单步串行修正极轻（论文 Figure 4 右图：γ=16 时延迟仅增 1.3%）。
4. 三项 loss 与 acceptance rate 直接对偶：CE/TV/BCE 都可解析推导到 $\tau$，训练目标与评测目标对齐。
5. confidence head 输入含 markov embedding：让置信度看到前一个 draft token，建模条件接受率。
6. world_size 缩放：loss 乘 world_size 匹配 all_reduce SUM 的分母聚合。

易踩坑点：

• aligned_target_logits 必须从 target_last_hidden_states 算——它用 draft 模型的 compute_logits（共享的 lm_head），所以需要在 prepare_target_cache 阶段额外存 target_last_hidden_states（每 token 1 个 hidden，非 5 层）。
• eval_mask 的 cumprod 在 fp32 上做，bf16 上数值不稳定。
• num_anchors 过大会 OOM——默认 512 在 8×A100 80G 上刚好。
• markov_head_type='rnn' 推理时需要维护状态，不能简单并行化，默认用 vanilla。
• confidence_head_with_markov=True 需要 markov_head 不为 None，即 markov_rank > 0——DFlash 配置下不能开。

延伸阅读：进入 04 Eagle3 建模 看自回归 drafter 怎么做对照；进入 07 评测系统 看 Qwen3DSparkEvaluator 如何在推理时调用 forward_dspark_draft_block 与 build_dspark_proposal；进入 08 实验复现 看如何用 --opts 调 block_size/markov_rank 复现论文 Figure 3、4。论文 Section 3.1（Markov head 公式 5、RNN head 公式 6）、3.2.1（confidence head 公式 7、8）、3.3（loss 公式 9-12）在 [DSpark_paper.pdf](file:///workspace/DSpark_paper.pdf)。