TreeSpark: Calibrated, Load-Adaptive Draft Trees for Semi-Autoregressive Speculative Decoding

Huapeng Zhou; Huayu Wang; Xinyu Wang

cs.CLAug 12, 2026

TreeSpark: Calibrated, Load-Adaptive Draft Trees for Semi-Autoregressive Speculative Decoding

Authors: Huapeng Zhou, Huayu Wang, Xinyu Wang

Organizations: Boson AI · University of Washington · McGill University

Abstract

Speculative decoding accelerates language-model inference by letting a cheap drafter propose tokens that the target model verifies in parallel. Recent block drafters make drafting nearly free: a single backbone pass emits an entire block of draft tokens. Draft trees promise a further gain -- several alternative continuations verified in one target forward -- but existing constructions rank candidates by per-position marginals that ignore which parent a candidate extends, so on semi-autoregressive drafters wider trees mostly add mis-ranked nodes; and a tree of fixed size ignores how much speculation each decoding round, and each serving load, can support. We introduce TreeSpark, which reads a parent-conditioned distribution from the drafter's existing Markov head at negligible cost, calibrates it into an edge-acceptance estimate, and lets path survival govern everything else: best-first expansion, per-round stopping, and a load-adaptive serving policy. Sampling siblings without replacement, with matching residuals in recursive rejection, keeps decoding lossless at any temperature. Adaptive trees improve on matched fixed budgets at every temperature; against a tuned chain on the same drafter, TreeSpark accepts 15-25% more draft tokens per round and decodes 8-14% faster in single-request wall-clock, and under rising load it gracefully shrinks the tree back to the chain. Code and artifacts: https://github.com/PopSoda2002/TreeSpark

Figures & tables

Figure 1: One round of TreeSpark. (1) A shared drafter backbone produces block logits, while the rank-

r

Markov head supplies a parent-specific conditional distribution. (2) Products of calibrated edge-acceptance probabilities estimate path survival, and best-first expansion stops when the best remaining candidate falls below

\theta

. (3) A single target forward verifies the resulting tree by exact matching at

T{=}0

or recursive rejection sampling at

T{>}0

. The runtime adjusts

\theta

as load changes, shrinking the tree toward the chain when verification becomes more expensive.

	$T{=}0$	$T{=}0.5$	$T{=}1.0$
$\text{bs}{=}1$	$0.02$ : ${+}7.5\%$ (5/5)	$0.05$ : ${+}6.1\%$ (3/3)	$0.05$ : ${+}3.0$ – $3.4\%$ (7/8)
$\text{bs}{=}16$	$0.15$ : ${+}2.1\%$ (4/4)	chain (parity)	chain (parity)

Table 1: Measured price ladder

\theta^{*}(\text{bs},T)

constructed by Algorithm 2 . Each strict-tree entry reports the selected

\theta

, its wall-clock margin over the chain, and its paired win count (A100 engine; same-GPU interleaved A/B runs with prebuilt engines and the warm-up repetition discarded). Each selected

\theta

also beat the neighboring probes. At

\text{bs}{=}16

and

T{>}0

, no branching tree was faster than the chain. At

T{=}1

, the probes

\theta\in\{0.05,0.10,0.15,0.4\}

trailed by

12/4.6/2.4/7\%

; at

T{=}0.5

\theta{=}0.15

trailed by

1.5\%

(0/5). The ladder therefore selects its chain member in those cells, giving parity by code identity. Intermediate batch sizes use the shipped interpolation stated below; Fig. 3 (a) provides modeled context rather than a cell-by-cell certification.

Figure 2: Round-by-round

\theta

adaptation under churn (48 requests, Poisson arrivals at

4

/s,

T{=}0

, A100). Left: The ladder reads the active batch each round. Trees shrink as the batch fills and grow again as it drains; near the end of the run, the controller returns to the

{\sim}25

-node regime associated with

\theta{=}0.02

. Right: Distribution of tree sizes at each selected price. Higher

\theta

shifts the distribution toward the chain, while variation within each price reflects the survival profile of the individual request.

Figure 3: Serving results at

T{=}0

; colors denote the same policy in both panels. (a) Combining measured step costs with measured acceptance shifts the best shape from the largest tree in the memory-bound regime (

\text{bs}{\leq}4

) to small adaptive trees in the compute-bound regime (

\text{bs}{\geq}16

); the adaptive curve is the modeled best

\theta

at each load. (b) In a scheduler-level simulation with request churn, each fixed shape saturates at a different rate. The load-adaptive tree reaches the highest capacity, about

{+}15\%

over the chain; the fixed

N{=}56

tree, which is optimal at

\text{bs}{=}1

, has the lowest capacity under load. Real-engine measurements in the text confirm the same trends.

Target	Method	$N{=}7$	$N{=}14$	$N{=}28$	$N{=}56$
Qwen3-4B	DFlash chain	3.86	–	–	–
	DDTree	4.18	4.60	4.89	5.13
	DSpark chain	4.39	–	–	–
	TreeSpark	4.57	5.18	5.50	5.81
Qwen3-8B	DFlash chain	3.96	–	–	–
	DDTree	4.25	4.66	4.98	5.19

Table 2: Accepted tokens per target forward (

\tau

) under greedy decoding on 150 prompts from the original six benchmarks. Budgets count verified draft tokens per round (

N

is a fixed per-round node quota; the

\theta

stop is not used here, so this isolates parent conditioning at budgets matched to DDTree — § 5.3 compares quota against price). Chains cannot exceed

N{=}7

: the block-7 drafter yields seven depths per backbone pass, and only branching spends more. Markov-conditioned trees improve monotonically with budget and outperform marginal DDTree at every matched budget. Finding (1) reports the marginal-scored DSpark ablation, which underperforms the chain.

Target	Config	Verified	$\tau$	AR speedup	Util.	Uncommitted
Qwen3-4B	chain $N{=}7$	7.00	4.961	3.124 $\times$	51.23%	48.77%
	TreeSpark $\theta{=}0.15$	6.96	5.179	3.198 $\times$	56.92%	43.08%
	TreeSpark $\theta{=}0.05$	12.99	5.795	3.373 $\times$	33.93%	66.07%
	TreeSpark $\theta{=}0.02$	23.17	6.150	3.318 $\times$	20.19%	79.81%
Qwen3-8B	chain $N{=}7$	7.00	5.183	3.158 $\times$	53.13%	46.87%
	TreeSpark $\theta{=}0.15$	7.02	5.356	3.153 $\times$	58.63%	41.37%

Table 3: Calibrated adaptive evaluation at

T{=}0

on the nine-task suite (153 held-out prompts per model, three H200 repetitions). Verified is the mean number of non-root draft nodes per round; AR is speedup over the matched target-only decoder; Util. is the fraction of verified nodes on the committed path; Uncommitted is

1-\text{Util.}

Bold marks the highest AR among displayed rows. Fixed-budget H200 rows are omitted for compactness; the best omitted AR is

3.463\times

for

N{=}14

on 14B, compared with

3.458\times

for TreeSpark

\theta{=}0.05

. Separate Qwen3-4B frontiers under the 102-prompt protocol appear in Table 5 and Fig. 4 . Per-dataset

\tau

on the three additional tasks appears in Table 7 (App. B ).

	$T{=}0$		$T{=}0.5$		$T{=}1.0$
Decoder	$\tau$	AR	$\tau$	AR	$\tau$	AR
DSpark chain $N{=}7$	5.07	3.30 $\times$	5.13	3.14 $\times$	4.97	3.06 $\times$
TreeSpark	6.34	3.75 $\times$	5.92	3.38 $\times$	5.72	3.38 $\times$

Table 4: DSpark chain and adaptive TreeSpark with the final recompute implementation (A100, six-task held-out suite, 102 prompts, one same-GPU sequential pass per temperature). AR is wall-clock speedup over the matched target-only decoder in the same cache-free reference harness (both arms cache-free, unlike the KV-cached H200 protocol of Table 3 ), which runs at

30.3

tok/s for

T{=}0

and

28.8

tok/s for

T{>}0

. The adaptive policy uses the batch-size-1 ladder value:

\theta{=}0.02

T{=}0

and

\theta{=}0.05

otherwise. These AR values are comparable only within this recompute harness. Earlier stored-proposal fixed-

N

ladders remain available in the artifacts.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Config	$T{=}0.5$	$T{=}1.0$
chain ( $N{=}7$ )	7.0 / 5.13	7.0 / 4.97
fixed $N{=}14$	13.6 / 5.87	14.0 / 5.59
fixed $N{=}28$	26.4 / 6.20	28.0 / 5.92
fixed $N{=}56$	51.1 / 6.52	56.0 / 6.28
TreeSpark $\theta{=}0.15$	7.1 / 5.36	7.5 / 5.08
TreeSpark $\theta{=}0.05$	13.2 / 5.93	14.6 / 5.73

Appendix

Table 5: Fixed and adaptive variants of the method in § 3.2 –§ 3.5 at

T>0

(Qwen3-4B, 102 held-out prompts). Each entry reports verified draft tokens per round /

\tau

. When calibration uses

q_{0}

, adaptive stopping lies above the matched fixed frontier at both temperatures. Bold identifies the adaptive member of each matched-cost comparison, not the column maximum.

Full pool (descriptive)
Acceptance predictor	AUC	ECE
confidence head $c$ (raw)	0.699	0.260
Markov conditional $q_{0}$ (raw)	0.883	0.131
Held-out odd-indexed prompts
$c$ $+$ Platt (all-edge fit)	0.702	0.070
$q_{0}$ (raw) †	0.948	0.042

Appendix

Table 6: Acceptance prediction on 34,216 Qwen3-4B tree edges. Full-pool rows are descriptive; held-out rows fit even-indexed prompts and evaluate odd-indexed prompts. † restricts evaluation to the 6,256 ancestors-accepted edges. The AUC change from

0.883

0.948

comes from this population change, not Platt scaling.

Figure 4: Acceptance–cost frontier (Qwen3-4B, 102 held-out prompts, greedy). Calibrated adaptive stopping (green) lies above the fixed-budget line (blue) at every operating point. The same stopping rule with the miscalibrated all-edges fit (gray, dashed) falls below it, with the largest gap at the most restrictive threshold.

Target	AIME25	LiveCodeBench	Arena-Hard-v2
Qwen3-4B	6.398	6.005	4.337
Qwen3-8B	6.756	6.279	4.460
Qwen3-14B	6.617	6.547	4.282

Appendix

Table 7: Per-dataset

\tau

on the three additional tasks for TreeSpark

\theta{=}0.05

in the three-scale evaluation.

bs	no-spec	chain $N{=}7$	fixed $N{=}56$	adaptive
1	31	140	182	176
4	123	563	745	683
8	247	1134	1034	1425
16	495	2034	1231	2153
32	1000	2555	1312	2718
64	1524	2798	1342	2984

Appendix

Table 8: Modeled

T{=}0

goodput (tokens/s) from measured step costs and acceptance. Adaptive is the per-load oracle envelope over

\theta\in\{0.02,0.05,0.15\}

, shown as an upper frontier rather than the shipped ladder at every intermediate load; bold marks the row maximum.

Estimator	TV to analytic $p$
Direct target sampling (reference)	0.066
TreeSpark sampled tree (ours)	0.047
Top- $k$ + naive rejection (control)	0.462

Appendix

Table 9: Empirical losslessness at

T{=}1.0

: total variation from the analytic target at a high-entropy position over 300 one-round trials. Direct target sampling is the finite-sample reference.

Speculative Decoding

Speculative decoding accelerates generation by drafting multiple tokens and verifying them in parallel with a larger target model.

See all 291 papers

TreeSpark: Learning to Draft the Trees You Verify for Speculative Decoding
Oct 14, 2026Jihao Zhao, Qianyu Jiang, Shuaishuai Zu +6Speculative Decoding LLM Inference Acceleration
School of Information, Renmin University of China · National Key Lab for Novel Software Technology, Nanjing University · Huawei Technologies Ltd.
Accelerating Speculative Decoding with Block Diffusion Draft Trees
Apr 14, 2026Liran Ringel, Yaniv RomanoSpeculative Decoding Parallel Speculative Decoding
Department of Computer Science, Technion – Israel Institute of Technology. · Department of Electrical and Computer Engineering, Technion – Israel Institute of Technology.
TreeGraft: Adaptive Multi-Drafter Grafting for Tree-Based Speculative Decoding
May 28, 2026Jiaming Fan, Daming Cao, Canchen Huang +6Speculative Decoding LLM Inference Acceleration
Key Laboratory of New Generation Artificial Intelligence Technology and Its Interdisciplinary Applications (Southeast University), Ministry of Education, China · Southeast University · Ant Group +1

Oct 14, 2026J/K move · Enter open · S save

TreeSpark: Learning to Draft the Trees You Verify for Speculative Decoding

Jihao Zhao, Qianyu Jiang, Shuaishuai Zu, Pengfei Li+5

School of Information, Renmin University of China · National Key Lab for Novel Software Technology, Nanjing University · Huawei Technologies Ltd.

Speculative Decoding LLM Inference Acceleration

Open paper AskPDF arXiv

Parallel drafters such as DFlash and DSpark reduce proposal latency, but early mismatches limit how much of each block can be accepted. Retaining alternatives introduces parent states and candidate choices that reference-conditioned training does not directly supervise. We introduce TreeSpark, a unified training and inference framework extending the semi-autoregressive DSpark architecture. It constructs parent-conditioned candidate trees from one backbone pass and replays their expansion, pruning, and budgeted packing during training. Survival negative log-likelihood weights retained edges by reference-derived path mass, while a capped coverage hinge promotes target-important alternatives near local selection boundaries. Both objectives reuse reference-chain target distributions without additional target forward passes. We evaluate greedy decoding with Qwen3-4B and Qwen3-8B on nine general benchmarks and six writing domains. Across targets and suites, TreeSpark improves mean acceptance length over DSpark by 23.4% to 31.6% and mean single-request throughput by 15.1% to 25.0%. Batched Qwen3-4B throughput improves by 12.6% to 17.2%. At fixed attention capacity and inference policy, the proposed supervision further improves mean acceptance length by approximately 6.4% and mean system throughput by 4.2% to 5.0% on Qwen3-4B.

Abstract

Speculative Decoding

Explore similar work

TreeSpark: Learning to Draft the Trees You Verify for Speculative Decoding

Accelerating Speculative Decoding with Block Diffusion Draft Trees

TreeGraft: Adaptive Multi-Drafter Grafting for Tree-Based Speculative Decoding