DDTree Deep Dive: Optimal Draft Tree Construction for Block-Diffusion Speculative Decoding

DDTree deep dive based on: *Accelerating Speculative Decoding with Block Diffusion Draft Trees* — Liran Ringel, Yaniv Romano (Technion), arXiv: 2604.12989 (2026.4.14).

• Code: https://github.com/liranringel/ddtree
• Project page: https://liranringel.github.io/ddtree/

Key points

1. The problem: what DFlash wastes

• DFlash's block diffusion drafter outputs per-position marginal distributions q₁,...,q_L over a masked block (length L, plus a bonus token b already chosen by the target model).
• Greedy DFlash takes the argmax at each position, yielding a single trajectory. Alternative tokens with substantial probability mass are discarded entirely.
• Crucially, the qᵢ are marginals, not conditionals — position *i* does not depend on earlier choices. DDTree exploits this by building a tree of the B most valuable prefixes instead of one path.
• Since target-model probabilities p(y₁:L|c,b) are unknown at draft time, DDTree uses the drafter's factorized Q(y₁:L) = ∏ᵢ qᵢ as a surrogate objective.

2. Mathematical core (three propositions)

• Notation: c = context, b = bonus token, L = block size, B = node budget, q(u) = ∏ᵢ qᵢ(uᵢ) = prefix probability under Q, T = prefix-closed draft tree.
• Proposition 1 — additivity: E_{Y~Q}[α_T(Y)] = Σ_{u∈T} q(u), where α_T(Y) = max{d : y₁:d ∈ T} is the acceptance length. Proof: write α as a sum of indicator functions 𝟙[y₁:d ∈ T], swap expectation and summation; each prefix contributes exactly q(u). No cross terms — the objective is additive.
• Proposition 2 — optimal structure: the top-B prefixes by q(u) form a valid (prefix-closed) tree and maximize expected acceptance length among all trees with ≤B nodes. Prefix-closure holds because a parent always has higher probability than its child. Tree construction reduces from combinatorial optimization to ranking.
• Lemma 1 — search-space reduction: an optimal tree exists using only each position's top-K tokens (K = min(B, |V|)), shrinking the space from O(|V|^L) to O(K^L).
• Proposition 3 — best-first heap: expand candidates via sibling (replace last token with next-ranked) and child (extend one depth with top-1) moves in a max-heap scored by Σ log q. Complexity: O(B log B), heap size O(B), and provably optimal — unlike beam search (O(B·L·K), no guarantee) or brute force (O(K^L)).

3. Implementation highlights (ddtree.py and friends)

• Pipeline stages: draft → tree_build (copy / heap / visibility) → tree_compile → verify → commit, each independently timed for profiling. DFlash baseline has only draft/verify/commit.
• Heap is implemented with GPU-parallel tensor ops rather than a Python heapq; outputs are tree_tokens [B, budget] and tree_indices [B, budget, 2] (position, parent position).
• The drafter (Qwen3DFlashAttention, RoPE + GQA, FlashAttention-2 mandatory) conditions on hidden states sampled from intermediate target-model layers, skipping roughly the last three layers, which carry less useful conditioning signal.
• benchmark.py uses block_size=1 DFlash as the baseline for fair comparison.
• Flash Attention is required for the drafter; the target verifier is forced to SDPA because of the custom tree attention mask.

4. Tree attention: ancestor-only mask

• Each tree node may attend only to: past context (via KV cache), the bonus token, its ancestors along its own path, and itself — never siblings or other branches, to avoid contaminating branch representations.
• Unlike Medusa's static trees, DDTree's tree is rebuilt every round, so the mask is recomputed dynamically each iteration.
• Mask construction walks parent pointers from tree_indices recursively (−1 denotes the bonus/root).

5. KV cache compaction

• After verification only the accepted path survives; tree tokens are interleaved in the cache, so simple slicing is insufficient.
• An optional C++/CUDA extension (compact_attention, compact_tail_inplace) performs in-place sparse compaction for performance; a Python fallback using advanced indexing preserves correctness if compilation fails.
• Without compaction the cache would grow by tree_budget tokens per round; with it, only accepted_length.

6. Verifier walk

• Greedy: from the bonus token, follow target-model choices through the tree; on mismatch, stop; the first unmatched target token becomes the next bonus token.
• With temperature > 0, standard speculative-decoding accept/reject rules preserve the target distribution exactly.

7. Comparisons with related work

• vs DART: DART uses parallel logits plus external N-gram/continuity heuristics; DDTree needs no external resources and carries theoretical guarantees (Props 1–3).
• vs OPT-Tree: OPT-Tree needs L drafter forward passes (one per depth); DDTree needs one — the core block-diffusion advantage.
• vs Medusa: Medusa requires trained prediction heads and a static tree; DDTree is training-free given a DFlash checkpoint and adapts the tree per round.
• vs EAGLE-3: EAGLE drafts autoregressively layer by layer; DDTree drafts in a single forward.

8. Experimental observations

• Tested on Qwen3-4B/8B and Qwen3-Coder-30B on 8× H200, with B ∈ {16,...,1024}, temperature 0.0 and 1.0, across Math (GSM8K, MATH-500, AIME24/25), Chat (Alpaca, MT-Bench), and Code (HumanEval, MBPP, LBPP, SWE-bench, LiveCodeBench).
• Largest gains (~60%+) on open-ended tasks (Alpaca, MT-Bench); moderate (~20–40%) on math/code; smallest (~10–20%) on highly deterministic tasks (AIME, SWE-bench Lite) — gains correlate with task diversity.
• B = 64–128 is the sweet spot; returns diminish beyond 256–512. Sampling (T=1.0) benefits more than greedy.

9. Future directions and tech debt

• Adaptive block size L and adaptive budget B (e.g., entropy-driven); joint fine-tuning with the target; extension to multimodal.
• Debt: CUDA-extension portability, Qwen3/GQA-specific tree attention, hard flash-attn dependency, batch-size-1 optimization focus, and memory peaks during tree building for large B.

10. Engineering philosophy

The whole algorithm is ~50 lines of pseudocode: one drafter pass → top-K per position → max-heap top-B prefixes → build mask → one target pass to verify. No extra trained components, no external data structures — minimalism as a design principle.