DFlash 与 DSpark：从原理到 SGLang 实现

DFlash 与 DSpark：从原理到 SGLang 实现

本文梳理 Flash-style block speculative decoding 里两条相关路线：DFlash （block diffusion + KV injection）与 DSpark（半自回归 refine + confidence 调度 verify），并对照 SGLang 中的落地方式。

---

1. 背景：Speculative Decoding 在优化什么

自回归解码每步只出一个 token，算力利用率低。Speculative decoding 用更轻的 draft 提出多个候选，再用 target 一次（或一块）校验：

L=Tdraft+Tverifyτ L = \frac{T_{\mathrm{draft}} + T_{\mathrm{verify}}}{\tau} L=τTdraft+Tverify

加速要么靠提高平均接受长度 (\tau)，要么降低 draft / verify 开销。

• EAGLE 类：draft 仍多为自回归，成本随 draft 长度近似线性增长。
• DFlash / DSpark ：偏向 一块 token 并行提出，更好吃满 GPU；同时用 target 的内部表示强化 draft 条件。

---

2. DFlash：在干什么

2.1 一句话

用 target 的 hidden 当 draft 的历史记忆（写入 draft KV）；draft 只在当前 MASK block 上并行猜词；再用 target 做 verify。

2.2 与「普通 attention / MLA」无关的一点澄清

DFlash 解决的是 投机解码的 draft--verify 管线，不是模型内部的 MLA（QK 小维、V 大维）那套。二者可同时存在于服务栈里，但概念层不要混。

2.3 KV Injection（相对 EAGLE 的核心差别）

EAGLE 常见做法：把 target 特征融进 draft 输入 ，draft 自己算完整 QKV，维护 私有 KV。层一深，target 信号容易变淡。

DFlash：把融合后的 target 上下文 直接注入每一层 draft 的 K/V cache。

Step 1 --- 抽多层 target hidden 并投影到 draft 隐空间（各层共享）：

Ht=RMSNorm(WcH(l1);...;H(lm)) \mathbf{H}_t = \mathrm{RMSNorm}\big(W_c\\mathbf{H}\^{(l_1)};\\ldots;\\mathbf{H}\^{(l_m)}\big) Ht=RMSNorm(WcH(l1);...;H(lm))

Step 2 --- 每层只写 K/V，不写 Q：

Qi=WiQHdKi=WiKHt; WiKHdseqVi=WiVHt; WiVHdseq \begin{aligned} \mathbf{Q}_i &= W_i^Q \mathbf{H}_d \\ \mathbf{K}i &= W_i\^K \\mathbf{H}_t;\\, W_i\^K \\mathbf{H}_d{\mathrm{seq}} \\ \mathbf{V}i &= W_i\^V \\mathbf{H}_t;\\, W_i\^V \\mathbf{H}_d{\mathrm{seq}} \end{aligned} QiKiVi=WiQHd=WiKHt;WiKHdseq=WiVHt;WiVHdseq

含义：

• Ht\mathbf{H}_tHt：历史上下文（prompt + 已接受 token）------来自 target。
• Hd\mathbf{H}_dHd：当前 draft block 自己的表示。
• Draft 的 Q 去查：「target 注入的 KV」+「block 内 KV」。

因此 draft 不必自己把整段 prompt 再建模一遍。

2.4 Block / MASK 长什么样

例如 block size = 4，上一个已接受 token 是 cat：

位置	输入	含义
0	cat	anchor（已知正确）
1--3	[MASK]	待并行填空

Draft 一次 forward 填出例如 sat on the，再交给 target verify。

Block 内 attention 多为 双向（diffusion 式），mask 位之间可以互相看，从而「一块一起猜」，而不是 EAGLE 那样一步一个。

2.5 Prefill / Decode 与两套 KV

DFlash 里 target KV 与 draft KV 分离：

	Target KV	Draft KV
谁写	target 正常 QKV	target hidden → (W_c) → 各层 (W^{K/V}) 注入
用途	下次 verify	下次 draft block

Prefill： target 全量 forward（抓 FULL hidden）→ 立刻 materialize 进 draft KV → 留下 bonus（第一个生成 token）当下一轮 anchor。

Decode 循环：

1. 组 [bonus, MASK, ...]，draft 一次出提案
2. Target verify（同样抓 hidden）
3. 按接受前缀 commit：只把通过前缀的 target hidden 注入 draft KV（不是整块、也不是 draft 自己的 hidden）
4. 新的 bonus + seq_len → 下一轮

「Commit 前缀 hidden」= target verify 算出的 hidden，按 commit_lens 写入 draft KV。被拒尾巴不写，避免污染上下文。

2.6 和 EAGLE：verify 都要不要算 hidden？

两边 verify 通常都要 target hidden，差别在用法：

	DFlash	EAGLE
Hidden 用途	注入 draft K/V	喂给 draft_extend 输入
Draft 建 cache	materialize（投影写 cache）	draft 再跑一遍 extend

---

3. DSpark：在 DFlash 思路上多两刀

直觉上可记：

DSpark ≈ block 并行 draft（同类 conditioning）+ 半自回归 Markov refine + 按 confidence 动态缩短每人 verify 长度

3.1 半自回归 refine（Markov head）

纯并行 draft 容易「后缀衰减」（块内 token 缺依赖）。DSpark 在并行 backbone 给出 base_logits 后，用轻量 Markov head 逐步修正再采样：

并行 draft forward → base_logits[1..γ]
        ↓
for t in 1..γ:                    # 半 AR
    logits_t = base_t + Markov(prev_token)
    token_t  = sample(logits_t)
    prev = token_t

相对 DFlash「mask 位 hidden 直接过 lm_head」：多了一段 块内串行依赖建模，draft 质量（尤其后缀）更好，代价是 draft 侧多一点逐步采样（远小于再跑一层大模型）。

3.2 Confidence 调度 verify 长度

高并发下，固定每人验满 (K) 个 token，batch 侧是 (B\times K) 的 target 成本；后面位置接受率往往很低，等于浪费 verify。

DSpark：

1. 用很小的 confidence head （本质是 Linear → 1 + sigmoid）估计各位置「前缀还能活」的置信度
2. 结合负载 / SPS cost table，得到总预算与 每请求 verify_lens
3. Ragged / compact verify：置信度低的尾巴直接不验

Confidence 与定 lens 的开销相对 一次 target verify 通常可忽略；换来的是少验高拒位。静态全长 verify 模式下甚至可以不启用 confidence head。

3.3 和 DFlash 的「代码级」主差异（就记这两条）

	DFlash	DSpark
Draft 出 token	并行 MASK → 直接采样	并行 logits + Markov refine
Verify 长度	全员固定 block_size	confidence → 每人不同 verify_lens

KV 注入、target verify、accept 前缀、commit 等主流程两边同族，实现拆分不同，但不是概念主差。

---

4. SGLang 里落在哪

4.1 启用方式

# DFlash
--speculative-algorithm DFLASH
--speculative-draft-model-path <draft>
--speculative-dflash-block-size 8   # 可选

# DSpark
--speculative-algorithm DSPARK
--speculative-dspark-block-size <gamma>   # 可选，可从 checkpoint 推断

4.2 DFlash 关键文件

角色	路径
Worker 主循环	srt/speculative/dflash_worker_v2.py → forward_batch_generation
Draft 模型 / (W_c) / kv_proj_only	srt/models/dflash.py
融合写 KV	srt/speculative/triton_ops/fused_kv_materialize.py
prepare_block / accept	srt/speculative/triton_ops/dflash.py
状态	dflash_info.py / dflash_info_v2.py

主循环阅读地图（dflash_worker_v2.py）：

1. Prefill（extend）：target FULL hidden → _append_target_hidden_to_draft_kv_by_loc → 产出 bonus
2. Draft：[bonus\|MASK...] → draft_model_runner.forward → 采样提案
3. Verify：target is_verify=True + accept / commit_lens
4. 再注入：仅 commit 前缀的 target hidden → 下一轮 draft input

4.3 DSpark 关键文件

角色	路径
Worker	srt/speculative/dspark_components/dspark_worker_v2.py
Draft / Markov 采样	dspark_draft.py + models/dspark.py（markov_head / confidence_head）
Verify 调度	dspark_planner.py（compute_confidence_tensor / resolve_verify_token_budget / schedule_layout）
Compact verify	dspark_verify.py 等

Decode 上可对照阅读：先 propose（含 Markov），再 confidence → schedule_layout → ragged verify。

4.4 建议阅读顺序

1. 只读 worker 主路径，画出 prefill → draft → verify → inject
2. 再读模型里的 projection / Markov / confidence
3. 最后看 Triton / planner 细节
4. 需要时拿 EAGLE 的 draft_extend 对照「输入融合 vs KV 注入」

---

5. 一张总图

                    ┌──────────── target ────────────┐
Prompt / 已接受 ──►│  forward（prefill / verify）    │
                    │  产出 hidden（多层 concat）      │
                    └─────────────┬──────────────────┘
                                  │ W_c + 各层 W^{K/V}
                                  ▼
                    ┌──────────── draft KV（历史）────┐
                    │  注入的上下文 K/V                 │
                    └─────────────┬──────────────────┘
                                  │
     [anchor, MASK...] ──► draft ───┼──► Q 查（历史 KV + block KV）
                                  │
                    DFlash: 直接采样提案
                    DSpark: + Markov 半 AR refine
                                  │
                                  ▼
                         target verify
                    DFlash: 固定窗长
                    DSpark: confidence → 变长 verify_lens
                                  │
                          接受前缀 commit
                    （target hidden → 再注入 draft KV）

---

6. 可带回家的三句话

1. DFlash 用 KV injection 让小 draft 借用 target 的上下文记忆，并用 MASK block 一次提出多 token。
2. Commit 的是 target 的 hidden → draft KV ，且只提交 接受前缀；两套 KV 并不共享底层 buffer。
3. DSpark 在同族管线上加上 Markov 半 AR refine 与 confidence 变长 verify，主打「草案更稳 + 高并发少验废 token」。

---

参考

• DFlash 论文与 LMSYS / SGLang 集成说明（KV injection、block diffusion）
• DSpark 论文与 LMSYS 博文（semi-AR + confidence-scheduled verification）
• SGLang 源码：dflash_worker_v2.py、dspark_components/*、models/dflash.py、models/dspark.py