DDTree Explained: From Block Diffusion to Optimal Diffusion Draft Trees for Speculative Decoding

Paper: *Accelerating Speculative Decoding with Block Diffusion Draft Trees* Authors: Liran Ringel, Yaniv Romano (Technion) arXiv: 2604.12989 (2026.4.14) Code: https://github.com/liranringel/ddtree

Key points

1. The problem: DFlash's waste

• DFlash uses block diffusion to generate marginal distributions for an entire token block in one forward pass, but verifies only a single trajectory.
• The per-position distributions qᵢ carry far more information than one sampled path uses; the rest of the probability mass is discarded.
• Core challenge: given a fixed node budget B, select the most valuable set of candidate paths from the per-position marginals.

2. Mathematical framework

• Given context c and bonus token b, the drafter's one pass yields L per-position distributions qᵢ(·|c,b), defining a factorized distribution Q(y₁:L|c,b) = ∏ᵢ qᵢ(yᵢ|c,b).
• The ideal objective (path-conditioned target probabilities) is infeasible; DDTree substitutes E_{Y~Q}[α_T(Y)] — maximizing expected acceptance length under the drafter's factorized distribution, over trees T with at most B nodes.

3. Key theorems

• Proposition 1 (decomposition): E_{Y~Q}[α_T(Y)] = Σ_{u∈T} q(u|c,b) — an additive sum over prefix probabilities.
• Proposition 2 (optimality): the B highest-probability prefixes automatically form a valid (prefix-closed) tree, which is the optimal solution.
• Lemma 1: considering only top-K tokens per position (K = min(B,|V|)) preserves optimality.

4. Best-first heap algorithm (Algorithm 1)

Instead of enumerating O(|V|^L) prefixes, index prefixes by rank tuple ρ = (ρ₁,...,ρ_d), where ρᵢ = k means position i takes its k-th most probable token. A max-heap starts at ρ = (1) (top-1 token everywhere); each pop pushes:
• sibling (ρ₁,...,ρ_{d-1}, ρ_d+1)
• child (ρ₁,...,ρ_d, 1)

Complexity: O(B log B) time, O(B) heap size.

5. Implementation highlights

• build_tree(draft_logits, tree_budget) returns tree_tokens [B] and tree_indices [B, 2] (node index, parent index) for the tree attention mask.
• Ancestor-only tree attention mask: each token attends only to the bonus token, its ancestors, and itself — no cross-branch interference during verification.
• KV cache management: after each round, compact_cache keeps only the accepted path's cache (C++ extension compact_attention.cpp for speed).
• Verifier walk: from the bonus token, apply the target model's decoding rule, match against tree children step by step; return the accepted path and a new bonus token.
• Repo layout: ddtree.py (core), dflash.py (base), model/ (drafter architecture, distributed support), benchmark.py, plotting/table scripts.

6. Comparison with related work

| Method | Drafter | Tree construction | Key difference | |---|---|---|---| | DFlash | Block diffusion | Single greedy path | Verifies one trajectory only | | DDTree | Block diffusion | Best-first heap from marginals | Single diffusion pass → optimal tree | | OPT-Tree | Autoregressive | Layer-by-layer forward + dynamic selection | One drafter forward per level | | DART | Parallel logits | N-gram pruning + trie | Needs external N-gram scoring | | EAGLE-3 | Autoregressive | Feature-based drafting | Multi-layer feature fusion |

DDTree advantages: single drafter pass, no external scorer, and a theoretical optimality guarantee under the surrogate objective.

7. Experiments

• Models: Qwen3-4B, Qwen3-8B, Qwen3-Coder-30B-A3B-Instruct; DFlash checkpoints (z-lab/dflash).
• Benchmarks: MATH-500, GSM8K, AIME 2024/2025, HumanEval, MBPP, LiveCodeBench, SWE-bench Lite, MT-Bench, Alpaca. Hardware: 8× H200; temperatures 0.0 and 1.0.
• Results: DDTree beats vanilla DFlash in all 60 settings (10 datasets × 3 models × 2 temperatures).
• Speedups vs autoregressive decoding: ~5-7x (4B), ~6-8x (8B), ~4-6x (30B MoE).
• Mean acceptance length τ (incl. bonus token): ~3-5 (DFlash) → ~4-7 (DDTree), a ~10-60% gain depending on budget and dataset.
• Budget tradeoff: B=16 for latency-sensitive use; B=64-128 is the sweet spot; B=256-512 shows diminishing returns.

8. Limitations

• Fixed block size L (e.g., 16); adaptive sizing is future work.
• DFlash needs 5 layers of target-model hidden states — memory grows with block size.
• Very long contexts still need optimization (sliding window only partially helps).
• The surrogate objective uses the drafter's factorized distribution, not the target's true path-conditioned distribution.
• Engineering notes: drafter requires FlashAttention (target can use sdpa); target model must support custom attention masks (e.g., Qwen GQA).

9. Why it matters

• Zero extra training: fully reuses existing DFlash checkpoints — plug and play.
• Single drafter pass means negligible drafting latency; tree construction costs only O(B log B); ancestor-only masking reuses existing tree attention machinery.
• Conceptually, DDTree shifts speculative decoding from single-path to multi-path, from heuristics to provable optimality, from bespoke to universal.

One-line takeaway: DDTree turns one block-diffusion pass's per-position distributions into an optimal draft tree via a best-first heap, pushing speculative decoding speedups from ~6x to 8x+ at zero additional training cost.

Suggested reading order

1. DFlash paper (arXiv 2602.06036) — block diffusion basics 2. DDTree paper Sections 3-4 — core algorithm 3. Code: start at build_tree() in ddtree.py 4. Run benchmark.py to reproduce Table 1 5. OPT-Tree and DART papers for the broader tree-based speculative decoding landscape