Accelerating Speculative Decoding with Block Diffusion Draft Trees

前置阅读:本文构建在 DFlash (块扩散 drafter) 的概念之上,如果对 DFlash 的 5 层 K/V 注入与单序列输出还有印象,直接读 §3 就行;否则建议先翻 07_DFlash 或 02_EAGLE_evolution_DFlash 的 DFlash 段。 关于 lossless 的 rejection sampling 论证,见 15_Lossless_Proof_Tutorial。

§1 动机:为什么 DFlash 丢掉树是一种浪费

1.1 历史脉络

回顾 spec decoding 这条线在"draft → tree"上的演化,可以画成两根并行的轨道:

• AR drafter 轨道: Vanilla SD (Leviathan / Chen) → SpecInfer (tree attention) → Medusa (multi-head) → EAGLE / EAGLE-2 / EAGLE-3 → OPT-Tree (在 EAGLE 等之上做动态预算分配)。这条线的共性: drafter 是 autoregressive 的, 一步生成一个 token 的 path-conditioned 分布 p(y_i | y_{1..i-1}); 想搞树就得多次前向, 或者像 EAGLE-2 那样在每一层都 expand。
• Block diffusion drafter 轨道: Block Diffusion (Arriola 2025) → DFlash (一次前向出整块,5 层 K/V 注入) → DART (用 N-gram trie 做 tree pruning)。这条线的卖点是一次前向就出 L 个 future positions 的 marginals,但 DFlash 只取 argmax 串成一条 sequence 拿去验证,把 marginal 分布的信息丢了。

DDTree 就是在这两条轨道的交叉点上写的:把 OPT-Tree 风格的"top-B prefix 选择"原理性地搬到 block diffusion 的单次 forward marginals 上,得到一个不需要多次前向的最优树构造算法。

1.2 别的方案为什么不够

1.3 为什么这事不平凡

表面上看,DFlash 已经一次前向出 L 个 marginals,要做树似乎只是"在每个位置 top-K 笛卡尔积一下"。但有三个非平凡问题:

1. 组合爆炸: 真要枚举所有长度 ≤ L 的 prefix, 有 ∑|V|^d ≈ |V|^L 个; 即便每位置截到 top-K, 还有 K^L 个 combination, 在 B = 16~1024 这个预算下还是太多。需要一个不枚举完全集的枚举方法。
2. 验证 topology: 树要让 verifier 一次 forward 就能跑完, 必须配合 SpecInfer 的 ancestor-only attention mask + tree-aware position ids; 否则 KV cache 就乱了。
3. Lossless 保留: 树验证的 lossless 性来自 verifier 用它自己的 decoding rule 沿着 drafted children 走一步; 不依赖 drafter 分布是不是 "正确的"。这一点恰好让 surrogate 选择没有 lossless 风险。
4. Surrogate 是否能优化: 真目标是 target model 的 expected acceptance length, 但 target 的 path-conditioned probs 我们拿不到 (要拿到就得跑 target, 等于解题要先求解); 必须用 drafter 的 factorized Q 做 surrogate, 还得证明这个 surrogate 上的最优解可以被简单算法 recover。

论文的优雅之处恰好是: 因为 drafter 分布 Q 是 factorized 的 (每个位置独立, 不带 path conditioning), prefix probability 自然分解成乘积, surrogate 目标分解成 prefix mass 之和,这个加性结构让"top-B prefix"成为最优解,而 best-first heap 又能严格按 prefix 概率降序枚举它们。

§2 背景速查

核心式子小抄

§3 关键观察: AR drafter 的分布 vs Diffusion drafter 的分布

这一节是整篇文章的立论根据。两类 drafter 给出的"分布对象"在数学上是不同的:

这件事的两面性

好处: drafter forward 次数恒为 1,无论 tree budget 是 16 还是 1024,drafter cost 不变。

代价: Q 是 factorized 的 — 它假设位置之间独立, 而 target p 不是。如果在 position 1 选了 token "积分", target 在 position 2 上的真实分布会被 strongly conditioned, 但 drafter 给的 q_2 看不到这件事。所以 Q 是真目标 p 的近似,DDTree 不在 Q 上回头去做 conditioning, 而是承认它就是 surrogate, 在 surrogate 上做 exact 优化。

反向论证: 如果有人想"再多跑几次 drafter, 把 path conditioning 补回来", 那就退化成 AR drafter 了 — block diffusion 的速度优势就消失。DDTree 的设计选择是: 承认 surrogate gap, 但拿 1 次 forward 的速度去换。

§4 Surrogate 目标与 Proposition 1/2

4.1 把目标改写成 prefix mass 之和

设树 T 包含若干 prefix u, 每个 prefix 的 q-概率为 q(u) = ∏ q_i(u_i)。论文 Proposition 1 把"期望接受长度"展开成按深度的指示函数和,然后用 disjoint events 论证:

关键步骤是: 在 fixed depth d, 事件 {Y_{1:d} ∈ T} 等于 ∪_{u 是深度 d 节点} {Y_{1:d} = u}, 这些事件互斥 (一个 sequence 只可能是一个 specific u),所以概率直接相加。Q 是 factorized 才让 Pr[Y_{1:d} = u] = q(u)。

4.2 解最优:top-B prefix 自动 prefix-closed

目标是 max_T ∑_{u∈T} q(u) s.t. |T| ≤ B 且 T prefix-closed。Proposition 2 的精彩之处:

1. 由于 q_i ∈ (0, 1), 任何严格后代 v 满足 q(v) = q(u) · ∏ q(...) < q(u)。即祖先严格大于后代。
2. 所以在 q 降序排列里, 祖先必然出现在后代之前。
3. 取 q 最大的 B 个 prefix, 自然包含每个被选 prefix 的所有祖先 ⇒ prefix-closed 自动成立, 不用额外约束。
4. 由于目标加性 + 非负 + 无约束 (除了 prefix-closure 已经被自动满足), top-B 就是最优。

4.3 Lemma 1: 只需要每深度的 top-K tokens

搜索空间还是大: 全空间 ∑|V|^d 在 |V| ≈ 150K 下天文数字。Lemma 1 证: 存在最优树 T* ⊆ S_K, 其中 S_K = {prefix 每位仅用各深度 top-K 的 tokens}, K = min(B, |V|)。直觉: 如果 T 中有一 prefix 在某位用了排名 > B 的 token, 那这个位置上排名 1..B 的 token 替换上去后, 都能得到更高 q 值的 prefix, 矛盾它已经在 top-B 中。

§5 Best-first heap 算法 (Algorithm 1)

5.1 用 rank tuple 做索引

关键技巧: 不存 token id, 存每位的排名。一个 prefix ρ = (ρ_1, ..., ρ_d) 表示"第 1 位选第 ρ_1 高概率的 token, 第 2 位选第 ρ_2 高的, ..."。这样邻居关系特别简单:

• Sibling (改最后一位排名 +1): (ρ_1, ..., ρ_d) → (ρ_1, ..., ρ_d + 1)。 score 增量: -log q^(ρ_d) + log q^(ρ_d+1)。
• First child (深一层, 在 d+1 位选 rank 1): (ρ_1, ..., ρ_d) → (ρ_1, ..., ρ_d, 1)。 score 增量: + log q^{(1)}_{d+1}。

5.2 算法骨架

heap H = {((1,), σ((1,)))}     # 初始: 第 1 位排名 1 的 prefix
T = ∅
while |T| < B and H 非空:
    pop ρ = (ρ_1,...,ρ_d) with max σ(ρ)
    T += prefix(ρ)
    if ρ_d + 1 ≤ K:
        push (ρ_1,...,ρ_d + 1) with score σ(ρ) − log q^(ρ_d) + log q^(ρ_d+1)
    if d < L:
        push (ρ_1,...,ρ_d, 1) with score σ(ρ) + log q^{(1)}_{d+1}
return T

5.3 为什么这个 heap 不会漏 / 不会重

论文 Proposition 3 的证明套路:

1. 每个 ρ 有唯一 predecessor pred(ρ): 若 ρ_d > 1, pred = (ρ_1, ..., ρ_d - 1); 若 ρ_d = 1 且 d > 1, pred = (ρ_1, ..., ρ_{d-1})。
2. 因此每个 ρ 只会被 push 一次 (作为它 predecessor 的 sibling 或 first-child),不会重复。
3. q(pred(ρ)) ≥ q(ρ) 永远成立: sibling case 是把 q^(ρ_d) 替成更大的 q^(ρ_d-1); child case 是去掉一个 < 1 的因子。
4. 从而 heap 在任意时刻总包含所有 unpopped tuple 的当前最大者: 沿任意 unpopped ρ 的 predecessor chain 往上找, 第一个 predecessor 已 popped 的祖先 ¯ρ 现在在 heap 里, 且 q(¯ρ) ≥ q(ρ)。所以下一次 pop 一定取到全局最大的 unpopped。
5. 归纳: 第 k 次 pop = 第 k 大的 prefix (在 S_K 内)。B 次后得到 top-B。

5.4 复杂度

最多 B 次 pop, 每次至多 push 2 个 ⇒ heap 总规模 O(B), 总操作 O(B log B)。跟 |V| 无关, 跟 L 无关。这意味着即便 vocab=150K, L=16, B=1024, 也只是几千次 heap 操作 + log factor, 几乎零开销。

§6 完整 worked example (vocab=4, L=3, B=6)

设小词表 V = {A, B, C, D}, drafter 一次前向给出 3 个位置的 marginals:

取 B = 6 ⇒ K = min(6, 4) = 4 (不剪)。我们用 rank tuple 跑算法:

得到树 T (6 个节点):

验证 walk 的具体追踪

把树打平送给 target:序列 [b, A, B, B, A, B, C], position ids = [0, 1, 1, 2, 2, 2, 3], attention mask 设为 ancestor-only。target 一次 forward 输出每个节点位置的 logits, 假设 target 的 greedy decoding 给出:

• 从 b 开始 → target argmax = "A" ✓ (匹配 depth 1 子节点 A) → 接受, walk 进入 A
• 在 A 节点 (target 已经为它 forward 过) → target argmax = "B" ✓ (匹配 [A,B]) → 接受, walk 进入 [A,B]
• 在 [A,B] 节点 → target argmax = "D" ✗ (树里 [A,B,C], 没 D) → walk 停下
• 结果: 接受 [A, B], 实际输出 [b, A, B, D]; 下一轮 bonus token = D。3 个 token 一轮(bonus + 2 accepted)。

对照 vanilla DFlash: 同样 marginals 下它会取 argmax-per-position = [A, B, C]; 但 target 第 3 步给的是 D, 所以也只接受 [A, B], 同样 3 token / 轮。在这个例子里 DDTree 没赢; 但如果 target 在 [A, ?] 时给的是 "A" (rank 2), DFlash 因为只验证了 [A,B,C] 这一条会在第 2 步就停下 (接受 1), 而 DDTree 因为同时挂了 [A,A] 这条分支, 还能继续接受 → 这是 DDTree 的结构性优势。

§7 验证: tree attention + verifier walk + lossless

7.1 tree attention

SpecInfer 的标准做法: 把树 flatten 成 1D 序列, 给每个节点配:

• position_id = 节点的 depth (root b = depth 0)
• attention_mask: 每个节点只看到自己 + 所有祖先 + past KV (已 prefill 的 context)

这样 target 一次 forward 就给出每个 drafted 节点上的 logits — 等价于"假设 walk 走到这里, target 的下一步分布是什么"。论文实验里只能用 PyTorch SDPA, 因为 FlashAttention-2 不支持这种 token-level mask。drafter 自己仍然可以用 FA-2 (它的 forward 不带这个 mask)。

7.2 verifier walk

从 root b 起, 每步用 target 的 decoding rule (greedy 或 temperature sampling) 选一个 token; 看它是不是当前节点在树里某个 child; 是 → 接受, walk 到该 child; 否 → 停下, 把这个 (不在树里的) token 当作 next bonus。 这一步不需要再 forward target, 因为所有节点的 logits 都在那一次 forward 里出来了。

7.3 是否 lossless?

关键: verifier 用它自己的 decoding rule, 跟 drafter 提议的具体分布无关。如果走到一个节点, target 的分布是 p, 它按 p sampling; sample 出的 token 恰好是 children 之一就接受, 否则停下用这个 token 作 next bonus。 整个过程的输出分布 = "在 target 的 p 下连续 sample 直到首次 mismatch" — 这跟 vanilla autoregressive 完全等价, 即和 drafter 选了什么 prefix 无关(drafter 选得好不好只影响速度, 不影响分布)。 因此 DDTree 是 lossless。 详见 15_Lossless_Proof_Tutorial。

论文用的 sampling 路线属于 SpecInfer "verifier 走自己的 sampling, 树作为 hint" 这一支, 不是 Leviathan 那种 importance-weighted rejection sampling。 这两种 lossless 路径在 spec decoding 里都成立, 区别在于前者要求 verifier 必须 sample 自己, 后者用 q(·) 提议 + 接受/拒绝。 DDTree 选前者是因为更天然适合树。

§8 实验关键结果

8.1 主表速读 (T=0.0)

读法: 加速增量在每个数据集 × 模型 × 温度组合 (10×3×2=60 项) 上都为正。τ (mean acceptance length, 含 bonus) 在 reasoning / code 上能从 7~8 涨到 ~10, 说明 DDTree 不仅在尾部分布上塞更多长 prefix, 也在中位数上系统性提升。Alpaca / MT-Bench 这种开放对话任务整体 τ 较低, 因为 target 自身的分布更分散, drafter 也更难命中。

8.2 budget-quality 曲线

8.3 acceptance length 直方图

论文 Figure 4 (这里只口述): vanilla DFlash 的 τ 分布峰在 5~7, 长尾在 13~16 较稀; DDTree (B=512) 的峰右移, 长度 ≤ 3 的比例显著降低, 长度 = 16 (full block accept) 显著上升。 这跟 §6 worked example 中"多挂一条 [A,A] 备选分支"导致更难全 reject 的直觉一致。

§9 与同类工作对比

立场与定位

• vs EAGLE-2/3: DDTree 不是"更好的 dynamic tree", 它是不同 drafter 类下的 tree 算法。 二者的核心区别不在树构造, 在 drafter forward 次数: DDTree 1 次, EAGLE-2 ~depth 次。 当 target model 越大, drafter 节省的相对越不重要; 当 drafter 占比高 (4B/8B target) 时, DDTree 优势放大。
• vs OPT-Tree: 算法骨架几乎同构 (top-B prefix + heap), 但 OPT-Tree 在 AR drafter 上要靠多次 forward 拿 path-conditioned probs; DDTree 一次 forward 拿 factorized probs, 等价计算量从 D 次 drafter forward 降到 1 次。
• vs DART: DART 也是从 block diffusion 一次前向出发, 但用外部 N-gram trie+continuity score 来 score prefix; DDTree 的卖点是 — 完全不需要外部信号, drafter 自己的 logits 就是 surrogate, 而且这个 surrogate 上的最优解可被 simple heap exact recover。
• vs DFlash: 直接的 strict superset。drafter / KV / training 完全不动, 只是 verify 阶段把 sequence 换成 tree。 论文整张表 60 个 cell 都是 DDTree 赢。

§10 局限 / 个人 take / 待验证问题

局限

• FlashAttention-2 不支持 tree mask: target 端只能用 PyTorch SDPA, 实际部署到 H200 时这是个真痛点; 用 FlashInfer / paged tree-attention 替代是工程化必需。
• Surrogate gap: 在 Q (factorized) 上最优, 不等于在 p (path-conditioned) 上最优。 实际跑下来 τ 长尾上 (Alpaca / MT-Bench) 收益相对小, 跟这个 gap 一致 — 当 target 高度依赖 conditioning 时, factorized prefix score 与真实 acceptance probability 偏离更大。
• L 是固定的: drafter 一次出 16 个 marginals, 接受长度 τ ≈ 10 时还有 6 个白扔; 把 block size 也调成动态 (按预测 confidence 决定生成多远) 是 next step。
• 预算 B 是手调超参: 论文按数据集 sweep 取最佳; 工业部署不可能每个 prompt 重 sweep。 是否能像 EAGLE-2 那样基于 confidence 在线决定 B 是开放问题。

个人 take

• 这是"正确的论文" — 一旦读到 DFlash 一次前向出整块, 任何接触过 OPT-Tree 的人都会立刻问 "为什么不构树"。 这篇文章把这个 obvious 但又不平凡的工程跑通了, 而且证明逻辑干净 (Prop 1/2/3 + Lemma 1, 4 个性质就够), 没有 hack。
• 真正的核心其实是 §3 那个对照: factorized Q vs path-conditioned p。 一旦想清楚这一点, 后面所有结论 (top-B 自动 closed, sibling/child 邻居关系) 都是 mechanical 的。 这种"概念澄清就解锁工程"的论文很值得作为 ML 写作模板。
• 下一步最有价值的 follow-up 在我看来不是"better tree", 而是"better surrogate": 能否在 drafter 上加一个轻量 head 给出近似的 path-conditioned 分布 (类似 EAGLE 的 hidden state 进 head), 而保持 1 次 forward — 这就在不引入新 forward 的前提下缩小 surrogate gap。

待验证问题

1. FlashInfer 是否已经支持 token-level tree attention mask? 如果支持, DDTree 在 30B-MoE 上的 8.2× 还能再涨多少?
2. budget-quality 曲线在不同 GPU (A100 / B200 / H200) 上的最优 B 漂移有多大? Appendix B 说"can shift across hardware",但没量化。
3. Alpaca / MT-Bench 上 τ 提升相对小,是 surrogate gap 的体现, 还是 drafter 本身在开放对话上 q_i 的熵就高 (top-1 概率低)? 这两者影响 DDTree 收益的方式不同。
4. 如果把 DFlash drafter 换成 PARD (AR-mimicking diffusion), DDTree 是否还成立? PARD 的"diffusion"是模拟, 实际给的是不是真 factorized marginals?
5. DART 论文的 N-gram continuity score 跟 q(u) 的差异在 ablation 上有多大? 是不是 DART 的开销其实没换来什么?
6. Lemma 1 假设 q_i ∈ (0,1), 当 logits 出现 −∞ (constrained decoding / banned tokens) 时, 算法还需要处理 ρ_d+1 不存在的边界 — 工程上是否要清洗 vocab 后再跑 heap?

Memory points