JetSpec++: Foresight Parallel Drafting with Self Refinements

Zhaoxiang Feng; Lanxiang Hu; Haoran Yuan; Yu-Yang Qian; Jihui Yang; Peng Zhao; Yibo Zhu; Tajana Rosing

cs.CLOct 14, 2026

JetSpec++: Foresight Parallel Drafting with Self Refinements

Authors: Zhaoxiang Feng, Lanxiang Hu, Haoran Yuan, Yu-Yang Qian, Jihui Yang, Peng Zhao, Yibo Zhu, Tajana Rosing,

Organizations: UC San Diego · UIUC · Nanjing University · Meta ·

Abstract

Speculative decoding's effectiveness is governed by two quantities: the acceptance rate of drafted tokens and the cost of producing them. We show the two objectives can be co-optimized with algorithm-system co-design, and present JetSpec++, which retains a strong learned draft model while hiding its cost almost entirely. JetSpec++ combines three components: intra-block self refinement, which trains the draft model to iteratively correct its own parallel predictions within a block of tokens; inter-block foresight drafting, which begins the next block before the current one is verified; and pipelined drafting and verification, which overlaps the two stages so drafting's marginal wall-clock cost approaches zero. In comparison with state-of-the-art speculative decoding methods, JetSpec++ attains a wall-clock speedup of up to $7.45\times$ over autoregressive (AR) decoding and $1.54\times$ over DSpark, with up to 1,361 tokens per second for a single request on GB200, and an average acceptance length of 8.94 across seven benchmarks, while preserving the target model's output distribution. Our code and models are available at https://github.com/hao-ai-lab/JetSpec.

Figures & tables

Figure 1: Non-pipelined drafting versus JetSpec++. (a) A draft model that reads the target’s hidden states drafts only after verification. (b) JetSpec++ runs foresight and Jacobi passes while block

t

is verified. Segment lengths are schematic.

Figure 2: Illustration of Foresight Drafting. (a) Foresight Drafting drafts the set

S

of selected outcomes (Eq. 7 ) in one batch with

J{+}1

passes. (b) At each target layer

\ell

that the draft model reads, the adapter attends to the verified prefix states

h^{\ell}_{<n}

and predicts the current-block states

\hat{h}^{\ell}_{n:n+\gamma}

. A decoder converts the top-layer hidden state predictions into the outcome distribution

P(k,a)

of Eq. 6 . (c) The speculation cache for block

t{+}1

contains outcome

k

, the draft tokens

d_{1},\dots,d_{k}

that outcome

k

accepts, the corresponding adapted states

\hat{h}

the draft model will read, and the decoded bonus token

a

as the anchor of block

t{+}1

. An outcome whose calibrated score (Appendix B.1 ) reaches the threshold

\lambda

will be pre-drafted as a

\gamma

-token block.

Temperature = 0
Method	Budget	GSM8K	MATH-500	AIME25	HumanEval	MBPP	LCB	MT-Bench
DFlash	16	$5.91$	$7.59$	$6.89$	$5.16$	$4.91$	$5.45$	$3.24$
DFlash	32	$5.18$	$6.67$	$6.07$	$4.53$	$4.31$	$4.70$	$2.95$
DSpark	16	$7.31$	$9.01$	$8.35$	$6.57$	$6.17$	$6.92$	$3.89$
DSpark	32	$7.06$	$8.66$	$7.73$	$6.25$	$5.85$	$6.38$	$3.77$
Domino	16	$6.91$	$8.74$	$8.09$	$6.06$	$5.90$	$6.27$	$3.74$

Table 1: Acceptance length

\tau

at verification budgets 16 and 32 for Qwen3-8B, in the setup of Section 5.1 . JetSpec++ runs

J_{1}

with one context refresh and anchor replacement (

J_{1}

R_{1}

) and

J_{3}

with three of each (

J_{3}

R_{3}

), and caches the

B=16

highest-scoring outcomes per round (Appendix B.1 ).

Qwen3.8-27B, temperature = 0
Method	Budget	GSM8K	MATH-500	AIME25	HumanEval	MBPP	LCB	MT-Bench
DSpark	16	$8.02$	$8.94$	$8.30$	$8.19$	$7.62$	$5.19$	$4.00$
DSpark	32	$7.57$	$8.49$	$7.64$	$7.13$	$6.96$	$4.84$	$3.83$
JetSpec++, $J_{1}$ - $R_{1}$	16	$8.71$	$9.53$	$8.58$	$9.18$	$8.41$	$5.80$	$4.30$
JetSpec++, $J_{1}$ - $R_{1}$	32	$9.68$	$11.10$	$9.82$	$10.72$	$9.42$	$5.94$	$4.40$
JetSpec++, $J_{3}$ - $R_{3}$	16	$10.05$	$10.83$	$10.15$	$10.72$	$9.80$	$6.93$	$4.91$

Table 2: Generalization to a larger target. Acceptance length

\tau

for Qwen3.8-27B as in Table 1 .

Model	Benchmark	Method	GPUs	tok/s (speedup)	$\tau$	hit	$\tau_{\text{hit}}$	$\tau_{\text{miss}}$	round	verify	exposed
Qwen3-8B	GSM8K	AR	1	182.3 (1.00 $\times$ )	1.00	N/A	N/A	N/A	5.49	5.45	0.04
		DSpark	1	744.3 (4.08 $\times$ )	7.34	N/A	N/A	N/A	9.92	6.16	3.76
		DFlash	1	620.1 (3.40 $\times$ )	5.87	N/A	N/A	N/A	9.52	6.11	3.42
		SSD	2	276.5 (1.52 $\times$ )	4.06	95.7%	4.12	3.23	14.75	6.25	8.51
		JetSpec++	2	1148.8 (6.30 $\times$ )	7.78	92.9%	8.20	3.45	6.84	6.23	0.61
	MATH-500	AR	1	182.7 (1.00 $\times$ )	1.00	N/A	N/A	N/A	5.47	5.44	0.04

Table 3: Throughput, acceptance length and round time at temperature 0 and budget 16 (5 for SSD, Appendix E ), on NVIDIA GB200 GPUs with CUDA graphs (Section 5.1 ). JetSpec++ runs

J_{1}

with one context refresh and anchor replacement for Qwen3-8B and

J_{3}

with three of each for Qwen3.8-27B. Speedup is over AR. Hit is the share

p_{\text{hit}}

of rounds whose outcome was pre-drafted, and

\tau_{\text{hit}}

and

\tau_{\text{miss}}

are the mean numbers of tokens committed on hit and miss rounds, all three over the rounds after the first of each request, which has no cache. Round and verify are the mean times per round in ms, and exposed is their difference.

		GSM8K				MATH-500
Inputs	Training blocks	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$
token (ours)	clean prefix (ours)	4.34	4.87	5.01	5.05	5.91	6.75	6.99	7.08
token (ours)	fully noised	4.36	4.37	4.37	4.37	6.02	6.01	6.02	6.01
mask	clean prefix	4.37	4.57	4.63	4.67	5.93	6.29	6.47	6.56
mask	fully noised	4.34	1.23	1.26	1.23	5.99	1.28	1.32	1.29

Table 4: Self-refinement of the draft model. Acceptance length

\tau

over refinement passes at temperature 0 for the two-epoch ablation draft model, trained with token or mask inputs, and on blocks with a clean prefix drawn by the prefix-length schedule of Appendix E or on fully noised blocks only. The target’s true states and no context refresh, 164 prompts per benchmark, 256 tokens. Paired per prompt, every other row trails ours at

J_{1}

J_{2}

and

J_{3}

by at least 8 standard errors.

$J$	$R$	Anchor repl.	GSM8K	MATH-500	AIME25	HumanEval	MBPP	LCB	MT-Bench
0	0	–	5.54	7.33	6.62	5.07	4.77	5.60	3.15
1	0	–	6.43	8.22	7.46	5.80	5.43	6.34	3.46
2	0	–	6.72	8.52	7.77	6.04	5.65	6.66	3.54
3	0	–	6.81	8.59	7.82	6.12	5.72	6.71	3.57
1	1	no	6.86	8.45	7.86	6.15	5.71	6.76	3.65
1	1	yes	7.98	9.63	8.97	7.23	6.73	7.59	4.38

Table 5: Gain of context refresh and anchor replacement. Acceptance length

\tau

at temperature 0 and budget 16 for the draft model and in the setting of Table 1 , with

J

passes after

J_{0}

and

R

context refreshes per round. A context refresh alone re-predicts the target states of the draft tokens and keeps every draft token. With anchor replacement, it writes the decoder’s token

\hat{a}

of Eq. 13 into the block.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

$J$	States	GSM8K	MATH-500	AIME25	HumanEval	MBPP	LCB	MT-Bench
0	stale	2.89	3.62	3.47	2.70	2.58	3.17	1.93
	predicted	5.56	7.31	6.89	4.98	4.68	5.93	2.98
	true	5.70	7.52	7.14	5.09	4.81	6.11	3.09
1	stale	3.36	4.17	4.00	3.13	2.97	3.64	2.11
	predicted	6.41	8.26	7.79	5.67	5.33	6.65	3.28
	true	6.59	8.55	8.16	5.85	5.49	6.89	3.40

Appendix

Table 6: Acceptance length

\tau

of the final draft model for Qwen3-8B at budget 16, temperature 0 and no context refresh, with three sources of the target states at unverified positions. Stale states copy the last verified state into every new position, the final adapter predicts the states, and true states exist only after verification. Every round drafts only the block of the verified outcome, so no round misses. Prompts as in Table 1 (first turn only for MT-Bench), with at most 2,048 generated tokens.

(a) Predicted states
	Error $\epsilon_{\ell}$					Cosine					Decoder
Layers	1	9	17	25	33	1	9	17	25	33	Top-1	KL
15 (2/2/2/4/5)	0.260	0.255	0.274	0.219	0.160	0.962	0.963	0.958	0.970	0.981	.958	.068
12 (1/1/1/4/5) (ours)	0.317	0.275	0.291	0.223	0.161	0.943	0.957	0.953	0.969	0.981	.956	.072
10 (1/1/1/3/4)	0.317	0.278	0.295	0.231	0.172	0.943	0.956	0.951	0.967	0.979	.952	.087
8 (1/1/1/2/3)	0.321	0.283	0.302	0.244	0.182	0.941	0.955	0.949	0.964	0.977	.945	.111

Appendix

Table 7: Effect of the number of adapter layers for Qwen3-8B. Parenthesized values indicate the layers per module for target layers 1, 9, 17, 25 and 33. We train each adapter for 0.46 epochs with

\kappa=1

in Eq. 10 . (a) Error

\epsilon_{\ell}

of Eq. 11 (lower is better) and cosine similarity between the predicted and the true states at each layer

\ell

that the draft model reads, on held-out mathematics. Top-1 is the fraction of block positions at which the decoder of Section 4.2 , reading each adapter’s predicted states, ranks the target’s greedy token first. KL is

\mathrm{KL}(p\,\|\,\hat{p})

in nats, where

p

is the target’s next-token distribution at temperature 1 and

\hat{p}

is the decoder’s. Both use the 200 held-out mathematics sequences of Table 9 . (b) Acceptance length

\tau

of the ablation draft model at temperature 0.

			GSM8K				MATH-500				Decoder
$\kappa$	$\epsilon_{25}$	$\epsilon_{33}$	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$	Top-1	KL
$0$	0.212	0.149	4.28	4.75	4.88	4.93	5.84	6.59	6.82	6.91	.967	.043
$1/4$	0.210	0.148	4.27	4.75	4.88	4.93	5.81	6.57	6.81	6.88	.967	.045
$1/2$ (ours)	0.208	0.147	4.28	4.76	4.89	4.95	5.86	6.59	6.82	6.89	.967	.043
$1$	0.209	0.148	4.27	4.75	4.87	4.92	5.82	6.55	6.80	6.88	.964	.050

Appendix

Table 8: Effect of the exponent

\kappa

of the position weight in the adapter loss (Eq. 10 ) for Qwen3-8B, with twelve-layer adapters trained for 1.83 epochs.

\epsilon_{25}

and

\epsilon_{33}

are the errors of Eq. 11 at layers 25 and 33 on held-out mathematics (lower is better), and the

J_{0}

J_{3}

columns give the acceptance length

\tau

of the ablation draft model at temperature 0. Top-1 and KL are as in Table 7 .

Figure 3: Scoring every outcome with one decoder pass. The adapter predicts the block’s states from the verified context, and the decoder reads them with the verified context as keys and values. Layers

\ell_{1},\dots,\ell_{\text{top}}

are the target layers that the draft model reads, and

\ell_{\text{top}}

is the deepest of them. Row

j

is the decoder’s

\hat{p}(\cdot\mid j)

. For the outcomes with

k=5

(outlined), rows 0 to 4 supply the acceptance probabilities

\hat{p}_{1},\dots,\hat{p}_{5}

of the draft tokens, row 5 supplies the bonus distribution with

d_{6}

excluded, and later rows are not read. The bars give each outcome’s score against the threshold

\lambda

of Section 4.2 : solid bars reach

\lambda

and are cached, and dashed bars are dropped. Heights are schematic. Our experiments instead cache the

B

outcomes with the highest calibrated score in every round (Appendix B.1 ).

Decoder	Top-1 agreement	KL to target
training-free	.967	.043
CE-trained	.968	.247

Appendix

Table 9: Offline agreement of each decoder with the greedy token of the Qwen3-8B target on 200 held-out mathematics sequences (379,305 positions), with the ablation adapter. KL is

\mathrm{KL}(p\,\|\,\hat{p})

in nats, where

p

is the target’s next-token distribution at temperature 1 and

\hat{p}

is the decoder’s.

	GSM8K				MATH-500
Decoder	$B=1$	$B=4$	$B=8$	$B=16$	$B=1$	$B=4$	$B=8$	$B=16$
training-free	.624	.927	.977	.993	.634	.903	.961	.989
CE-trained	.617	.934	.979	.994	.654	.935	.980	.995

Appendix

Table 10: Offline hit rate of each decoder when the cache holds at most

B

outcomes per round, ranked by the calibrated score of that decoder. The rounds come from JetSpec++ at

J_{1}

R_{1}

for Qwen3-8B with the final adapter in the setting of Table 1 , at temperature 0 and budget 16. The hit rate is the share of rounds after the first of each request whose verified outcome is in the cache.

	GSM8K				MATH-500
Rule	$B=1$	$B=4$	$B=8$	$B=16$	$B=1$	$B=4$	$B=8$	$B=16$
served rule	.624	.929	.978	.993	.644	.909	.963	.990
uncalibrated rule	.613	.921	.975	.993	.585	.870	.944	.981
draft-model rule	.387	.743	.885	.962	.444	.738	.862	.941
geometric fan-out	.050	.208	.512	.670	.044	.163	.588	.741
geometric fan-out, decoder anchors	.067	.289	.655	.890	.059	.222	.682	.898

Appendix

Table 11: Offline hit rate of each selection rule when it caches at most

B

outcomes per round, counted as in Table 10 and on the same rounds.

	GSM8K				MATH-500
Draft model	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$
JetSpec	5.86	2.39	2.65	2.45	7.51	2.57	2.84	2.61
JetSpec++	5.70	6.59	6.88	6.97	7.52	8.55	8.86	8.98

Appendix

Table 12: Acceptance length

\tau

over refinement passes for the final draft models of Qwen3-8B. All runs use the target’s true states in the setting of Table 6 . Refinement passes lower the

\tau

of the JetSpec draft model and raise that of the JetSpec++ draft model.

	GSM8K				MATH-500
$\Pr[m=0]$	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$
0	3.91	4.56	4.75	4.80	5.43	6.39	6.72	6.81
0.125	4.29	4.84	5.00	5.04	5.88	6.69	6.97	7.09
0.25 (ours)	4.34	4.87	5.01	5.05	5.91	6.75	6.99	7.08
0.375	4.32	4.78	4.90	4.94	5.91	6.69	6.91	6.97
0.5	4.41	4.84	4.96	5.01	6.01	6.74	6.95	7.03

Appendix

Table 13: Effect of the share of fully noised training blocks. Acceptance length

\tau

for Qwen3-8B in the setting of Table 4 . Each draft model takes token inputs and differs from ours only in the probability of prefix length

m=0

, with the rest split equally over lengths 1 to 15.

	GSM8K				MATH-500
Loss	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$
forward KL (ours)	4.34	4.87	5.01	5.05	5.91	6.75	6.99	7.08
SFT	4.28	4.74	4.85	4.88	5.81	6.58	6.76	6.81
reverse KL	2.94	3.06	3.08	3.08	4.03	4.29	4.34	4.36

Appendix

Table 14: Effect of the training loss. Acceptance length

\tau

for Qwen3-8B in the setting of Table 4 .

	GSM8K				MATH-500
Attention	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$	$J_{0}$	$J_{1}$	$J_{2}$	$J_{3}$
causal (ours)	4.34	4.87	5.01	5.05	5.91	6.75	6.99	7.08
bidirectional	4.36	4.81	4.95	5.00	5.88	6.65	6.82	6.93

Appendix

Table 15: Effect of the attention within the draft block. Acceptance length

\tau

for Qwen3-8B in the setting of Table 4 .

Figure 4: Correcting the first disagreement with one anchor replacement. The figure shows the context refresh after

J_{0}

of one batch row, the block

t+1

drafted for the outcome

(k,a)

, in the layout that Figure 3 uses for the predicted states of block

t

. The adapter re-predicts the states of block

t+1

, and the decoder reads them as in Figure 3 . The decoder’s top-1 token agrees with the

J_{0}

draft at positions 1 to 3 and disagrees at position 4, so

\hat{k}=3

and anchor replacement writes

\hat{a}

at position 4. The

J_{1}

pass reads

\hat{a}

at position 4 and the

J_{0}

tokens at the other positions. Anchor replacement corrects position 4 and fixes positions 1 to 4, and the

J_{1}

pass corrects position 7 by itself.

Method	Drafting	GSM8K	MATH-500
DFlash	critical path	5.91	7.59
DSpark	critical path	7.31	9.01
JetSpec++, $J_{1}$ - $R_{1}$	pipelined	7.98	9.63
JetSpec++, $J_{1}$ - $R_{1}$	critical path	8.10	9.82
JetSpec++, $J_{2}$ - $R_{2}$	pipelined	9.23	10.68
JetSpec++, $J_{2}$ - $R_{2}$	critical path	9.25	10.61

Appendix

Table 16: Acceptance length

\tau

with drafting pipelined or on the critical path, for Qwen3-8B in the setting of Table 1 at budget 16 and temperature 0. A pipelined draft starts before verification ends, as JetSpec++ serves. A critical-path draft starts after verification and reads the target’s true states, as in DFlash and DSpark.

Parallel Speculative Decoding

See all 28 papers

JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting
Jun 16, 2026Lanxiang Hu, Zhaoxiang Feng, Yulun Wu +9Speculative Decoding LLM Inference Acceleration
UC San Diego · Zhejiang University · UIUC +2
Speculative Speculative Decoding
Mar 3, 2026Tanishq Kumar, Tri Dao, Avner MaySpeculative Decoding LLM Inference Acceleration
Department of Computer Science, Stanford University · Department of Computer Science, Princeton University · Together AI
LibraSpec: Dynamic Diffusion-Based Speculative Decoding via Marginal-Gain-Driven Optimization
Aug 9, 2026Zexun Lin, Yuan Feng, Junlin Lv +2Inference-Time Optimization Speculative Decoding for Diffusion Language Models
Suzhou Institute for Advanced Research, University of Science and Technology of China Suzhou, Jiangsu, China

Jun 16, 2026J/K move · Enter open · S save

JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting

Lanxiang Hu, Zhaoxiang Feng, Yulun Wu, Haoran Yuan+8

UC San Diego · Zhejiang University · UIUC+2

Speculative Decoding LLM Inference Acceleration

Open paper AskPDF arXiv

Speculative decoding (SD) accelerates autoregressive Large Language Models (LLMs) by drafting multiple tokens and verifying them in parallel, but it faces a scaling limitation: increasing the draft budget improves speed only when acceptance remains high and drafting overhead stays low. This ceiling has been difficult to break because prior head-based SD methods face a causality-efficiency dilemma. Autoregressive drafters produce path-conditioned candidates that are effective for tree speculative decoding with higher acceptance length, but their drafting cost grows with tree depth. Bidirectional block-diffusion drafters generate all positions in one pass, but their branch-agnostic marginals can form individually plausible yet mutually inconsistent trees, wasting budget and reducing acceptance. We propose JetSpec, a head-based SD framework that combines one-forward drafting efficiency with branch-wise causal conditioning. JetSpec trains a causal parallel draft head over fused hidden states from the frozen target model, producing candidate trees whose scores align with the target model's autoregressive factorization. This enables JetSpec to convert larger draft budgets into longer accepted prefixes and higher end-to-end speedup. Across math, coding, and chat benchmarks on dense and MoE Qwen3 models, JetSpec consistently outperforms bidirectional-head and tree-based SD baselines. On H100 GPUs, JetSpec achieves up to 9.64x speedup on MATH-500 and 4.58x on open-ended conversational workloads, with further latency gains demonstrated through vLLM integration under realistic serving loads. Our code and models are available at https://github.com/hao-ai-lab/JetSpec.

Abstract

Parallel Speculative Decoding

Explore similar work

JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting

Speculative Speculative Decoding

LibraSpec: Dynamic Diffusion-Based Speculative Decoding via Marginal-Gain-Driven Optimization