PaperlayerView exact revision on arXiv

Source-linked AI summary

JetSpec++: Foresight Parallel Drafting with Self Refinements

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

arXiv:2610.19792v2cs.CL

TL;DR

Speculative decoding must balance acceptance rate against drafting cost, and cheaper model-free approaches can have high variance and low average acceptance. JetSpec++ overlaps foresight drafting with verification and trains iterative draft refinement. On GB200, it reaches speedups of up to 7.45× over AR and 1.54× over DSpark, with an average acceptance length of 8.94 across seven benchmarks.

  • Problem

    Speculative decoding must balance acceptance rate against drafting cost; the challenge is retaining model-based acceptance while removing drafting cost from the critical path.

    Source: §1 · p. 1

  • Method

    JetSpec++ overlaps target-conditioned drafting of the next block with current-block verification, using foresight predictions and Jacobi-style self-refinement with context refresh and anchor replacement.

    Source: §1 · pp. 1–2, §4.3 · p. 7

  • Results

    On GB200, JetSpec++ reaches up to 7.45× speedup over AR and 1.54× over DSpark; across seven benchmarks, its average acceptance length is 8.94.

    Source: §5.2 · p. 9, §6 · p. 11

  • Takeaways & Limitations

    Self-refinement and anchor replacement change only the proposal, so verification keeps the output lossless.

    Source: §4.3 · p. 7

    Serving integration remains improvable: fused kernels and optimized communication could reduce drafting overhead and enable more refinement passes and context refreshes.

    Source: §5.6 · p. 11

Abstract

from arXiv · show

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.

1 INTRODUCTION

JetSpec++ targets the trade-off between draft acceptance and drafting cost by overlapping target-conditioned drafting with verification. It predicts missing future context for the next block and uses concurrent compute for iterative self-refinement.

1 INTRODUCTION

  • Model-free drafting can be cheap but has high variance and low average acceptance, especially when reusable patterns are sparse or fail to generalize.JetSpec++ seeks to retain learned-model acceptance while removing drafting cost from the draft-then-verify critical path.Source: §1 · p. 1
  • Foresight Drafting anticipates possible acceptance-length outcomes, predicts corresponding target hidden states and anchor tokens, and uses them to draft the next block before verification finishes.This addresses the unavailable accepted-token count, bonus token, and target hidden states for the newly verified context.Source: §1 · p. 2
  • Jacobi-style self-refinement lets the draft model revise its parallel predictions over multiple passes while verification runs, using proxy verification, context refresh, and anchor replacement between passes.The extra draft computation adds no latency while it fits within the verification window.Source: §1 · p. 2
  • Up to 7.45× speedup over AR and 1.54× over DSpark, with 1,361 tokens per second for a single request and an average acceptance length of 8.94 across seven benchmarks.The reported speedups are end-to-end results on GB200.Source: §1 · p. 2

2 RELATED WORK

Related work addresses speculative decoding’s acceptance and drafting costs through tree verification, distillation, parallel drafting, and pipelining. JetSpec++ combines draft-model refinement with direct prediction of missing target hidden states for pipelining.

2 RELATED WORK

  • Tree verification scores multiple candidate branches to raise acceptance, while distillation aligns draft and target models; parallel drafting reduces draft cost by predicting a block in one pass.The cited approaches include Miao et al. (2024), Chen et al. (2024), and Li et al. (2024b) for tree verification; Zhou et al. (2024), Liu et al. (2024), and Goel et al. (2024) for distillation; and An et al. (2026) and Xiao et al. (2024) for parallel drafting.Source: §2 · pp. 2–3
  • Parallel-decoding methods include Jacobi fixed-point iteration and discrete diffusion, while JetSpec++ applies Jacobi iteration to the draft model rather than the target.Related approaches also train autoregressive models to reach fixed points in fewer iterations or accelerate diffusion inference.Source: §2 · p. 3
  • Prior pipelining methods hedge against unknown verification outcomes, but target-conditioned drafts lack the target representation of the context still under verification.JetSpec++ predicts the missing target hidden states directly, allowing a target-conditioned parallel draft model to speculate ahead without executing the target model.Source: §2 · p. 3

3 PRELIMINARIES

Speculative decoding verifies drafted tokens in one target-model pass, preserving the target distribution while trading draft cost against acceptance length. Parallel drafting reduces sequential draft work, and Jacobi iteration provides a multi-pass route toward autoregressive block predictions.

3 PRELIMINARIES

  • Speculative decoding verifies a drafted block in one target pass and losslessly commits the accepted prefix plus a bonus token; per-token latency is (Tdraft + Tverify)/τ.Here τ is the expected number of committed tokens per round.Source: §3 · p. 3

3.1 PARALLEL DRAFTING

  • Modern parallel draft models condition on target hidden states and predict masked block positions in one draft-model forward pass, unlike autoregressive drafts that require γ sequential passes.JetSpec makes the draft causal within the block, and JetSpec++ builds on that architecture for multi-pass Jacobi refinement.Source: §3.1 · p. 4

3.2 JACOBI DECODING FOR PARALLEL DRAFTING

  • Jacobi decoding updates every position from the previous iteration; under a causal mask, each iteration takes one forward pass and converges to the autoregressive fixed point under greedy decoding.JetSpec++ trains its draft model on both fully noised inputs and partially converged blocks.Source: §3.2 · p. 4

4 METHODS

JetSpec++ overlaps drafting the next block with verification of the current block, predicting the missing verification outcomes and hidden states so drafting can begin early. It uses intra-block self-refinement to improve draft alignment within the time pipelining frees.

4.1 DRAFT-VERIFICATION PIPELINING

  • Pipelining drafts block t+1 while block t is verified, caching blocks for anticipated outcomes and preserving the target output distribution through rejection-sampling verification.On cache hits, critical-path overhead is max(Tverify, Tdraft) + Tcomm; a miss adds fallback drafting time.Source: §4.1 · pp. 4–5

4.2 INTER-BLOCK FORESIGHT DRAFTING

  • Foresight Drafting predicts target hidden states and outcome distributions, then selects and batches likely outcomes so the next block can be drafted before verification finishes.The adapter predicts future target states, while the decoder estimates outcome probabilities; calibration selects candidate outcomes, subject to cache size B.Source: §4.1 · p. 4, §4.2 · pp. 6–7

4.3 INTRA-BLOCK SELF REFINEMENTS

  • Jacobi iterations let the causal draft model revise its own predictions, conditioning each position on earlier drafted tokens.The method adds context refreshes and anchor replacement between passes; verification remains lossless because these operations alter only the proposal.Source: §4.3 · p. 7
  • Training pairs clean prefixes with noisy suffixes and uses forward KL to train the draft model to match the target at every position.The clean-prefix length is sampled from a fixed schedule so the model learns to repair blocks with varying clean-prefix lengths.Source: §4.3 · pp. 7–8

5 EXPERIMENTS

JetSpec++ combines learned-draft refinement with foresight and pipelining to improve acceptance length and hide drafting latency. Across the evaluated targets and benchmarks, it outperforms the reported baselines in acceptance length and serving speed, while ablations identify contributions from refinement, context refresh, anchor replacement, and the adapter.

5.1 SETUP

  • The evaluation compares JetSpec++ with five baselines on Qwen3-8B and Qwen3.8-27B across seven benchmarks, using fixed-size blocks and up to 2,048 generated tokens.Serving is measured at batch size 1 on NVIDIA GB200 GPUs with CUDA graphs; speedup is reported relative to autoregressive decoding on one same-type GPU.Source: §5.1 · pp. 8–9

5.3 GENERALIZATION TO A LARGER TARGET

  • 8.71 vs. 8.02 on GSM8K and 9.53 vs. 8.94 on MATH-500: JetSpec++ J1-R1 leads DSpark on all seven benchmarks for Qwen3.8-27B at temperature 0 and budget 16.JetSpec++ also gains τ on all seven benchmarks at budget 32, whereas DSpark loses it.Source: Table 2 · p. 9, §5.3 · p. 9

5.2 RESULTS

  • At temperature 0 and budget 16, J3-R3 averages an acceptance length τ of 8.94 across seven benchmarks, versus 6.89 for DSpark with J1-R1 at 7.50.JetSpec++ also gains τ on every benchmark at budget 32, while DFlash, DSpark, and Domino lose it; on MATH-500, J1-R1 rises from 9.63 to 11.45 while DSpark falls from 9.01 to 8.66.Source: §5.2 · p. 9

5.4 SYSTEM PERFORMANCE

  • 1.54× faster than DSpark on GSM8K and 1.51× on MATH-500: JetSpec++ exposes 0.61 and 0.66 ms of drafting per round, versus DSpark’s 3.76 and 3.80 ms.The table reports throughput, acceptance length, round time, and verification time under CUDA-graph serving on GB200 GPUs.Source: Table 3 · p. 10, §5.4 · p. 10

5.5 ABLATION STUDIES

  • At J3, JetSpec++’s draft rises from 5.70 to 6.97 on GSM8K and from 7.52 to 8.98 on MATH-500, while the JetSpec draft falls on both benchmarks.The ablation indicates that token inputs and clean-prefix training blocks determine how much the draft model refines its predictions.Source: §5.5 · p. 11
  • Context refresh with anchor replacement adds 0.92–1.55 tokens at J1-R1 and 1.72–2.97 at J3-R3 over refinement alone; refresh alone adds at most 0.84.The reported refresh-with-replacement gains are 0.92 to 1.55 tokens at J1-R1 and 1.72 to 2.97 at J3-R3.Source: §5.5 · p. 11, Table 5 · p. 11
  • The adapter recovers 95% of the true-state gap; forward KL and causal attention outperform their tested alternatives.Forward KL leads SFT and reverse KL in every cell, by up to 0.27 and 2.72 tokens; causal attention leads bidirectional attention by 0.15 tokens on MATH-500 at J3.Source: §5.5 · p. 11

5.6 LIMITATIONS

  • Fused kernels and optimized communication could reduce drafting overhead and enable additional refinement passes and context refreshes to improve throughput.The authors state that the serving integration can be further improved through systems optimization.Source: §5.6 · p. 11

6 CONCLUSION

JetSpec++ drafts the next block while the current block is verified, using foresight predictions and self-refinement. On GB200 GPUs, it decodes up to 7.45× faster than autoregressive decoding and 1.54× faster than DSpark.

6 CONCLUSION

  • On GB200 GPUs, JetSpec++ decodes up to 7.45× faster than autoregressive decoding and 1.54× faster than DSpark.It drafts the next block during verification and uses Jacobi passes and context refreshes for self-refinement.Source: §6 · p. 11

A ADAPTER DESIGN AND ABLATIONS

The adapter predicts target hidden states at unverified positions, supporting the draft model’s acceptance length. The section also describes its architecture and training, then the decoder and calibrated outcome-scoring rule.

A ADAPTER DESIGN AND ABLATIONS

  • Each adapter module predicts one target layer’s block states from the preceding layer’s states, forming a chain from token embeddings to the deepest layer read by the draft model.The adapter has 2.32B trainable parameters for Qwen3-8B and 4.47B for Qwen3.8-27B.Source: A ADAPTER DESIGN AND ABLATIONS · p. 18
  • The adapter is trained against the frozen target on 16-position blocks, weighting positions by their probability of being reached; κ = 1/2 is used.Training uses 800k sequences less 24k held out, capped at 3,072 tokens.Source: A ADAPTER DESIGN AND ABLATIONS · p. 18
  • Predicted states yield 95 to 98% of true-state acceptance length across J0 to J3, and close 89 to 95% of the stale-to-true gap over 28 cells.The adapters use twelve layers, allocating extra layers to the deeper modules; reducing their layers lowers decoder agreement.Source: A ADAPTER DESIGN AND ABLATIONS · pp. 18–19

B.1 SCORING RULE AND CACHE SIZE

  • The decoder reuses frozen target layers above the deepest adapter-predicted layer, and calibrated outcome scores combine estimated acceptance length with a temperature-adjusted bonus distribution.The cache keeps the B highest-scoring outcomes per round; fitted scoring components include six coefficients, anchor temperature, and split-conformal thresholds.Source: §B.1 · pp. 20–21, §B.2 · p. 21

B.3 TRAINING-FREE VERSUS FINE-TUNED DECODERS

Fine-tuning the reused decoder layers yields little improvement in agreement and raises KL divergence. Among tested selection rules, the calibrated served rule leads or ties across cache sizes.

B.3 TRAINING-FREE VERSUS FINE-TUNED DECODERS

  • Fine-tuning raises top-1 agreement by .001 but increases KL divergence to the target distribution from .043 to .247.Its cache hit-rate improvement is at most .032, so the training-free decoder is served.Source: §B.3 · p. 22

B.4 COMPARISON OF SELECTION RULES

  • The served score combines calibrated outcome probability with the decoder’s bonus-token distribution, while comparison rules remove or replace sources of scoring information.The draft-model rule differs in its score of acceptance length; all rules rank anchors by the decoder.Source: §B.4 · p. 22
  • At B = 4, the calibrated served rule leads the uncalibrated rule by .008 on GSM8K and .039 on MATH-500, and leads or ties every rule at each cache size.The geometric fan-out trails the served rule by .640 on GSM8K and .687 on MATH-500 at B = 4, even with decoder anchors.Source: §B.4 · p. 22, Table 11 · p. 23

C TRAINING THE DRAFT MODEL TO REFINE

The draft model uses token inputs and clean-prefix training to refine its own predictions. Ablations show that fully noised blocks, loss choice, and within-block attention affect acceptance length and refinement.

C TRAINING THE DRAFT MODEL TO REFINE

  • A draft with mask inputs trained only on fully noised blocks collapses at J1, with acceptance length falling from 4.34 to 1.23 on GSM8K and from 5.99 to 1.28 on MATH-500.The final JetSpec++ draft instead uses token inputs and blocks with a clean prefix.Source: C TRAINING THE DRAFT MODEL TO REFINE · p. 23
  • At a fully noised-block share of 0.25, clean prefixes remain available for refinement; a share of 1 removes them, and the draft model does not refine.Shares of 0 and 1 produce lower J3 acceptance length than the four tested intermediate shares.Source: C TRAINING THE DRAFT MODEL TO REFINE · p. 23
  • Forward KL leads SFT and reverse KL, matching the target’s full next-token distribution rather than supervising only its greedy token.The gap between SFT and forward KL is larger at J3 than at J0.Source: C TRAINING THE DRAFT MODEL TO REFINE · p. 23
  • Causal attention matches bidirectional attention at J0, but the causal draft model leads during refinement passes.The comparison changes only the attention within the draft block.Source: C TRAINING THE DRAFT MODEL TO REFINE · p. 24, Table 15 · p. 24

D SELF-REFINEMENT EXAMPLE AND DRAFTING WITHOUT PIPELINING

Context refresh and anchor replacement let self-refinement repair draft positions, while pipelining keeps acceptance length close to critical-path drafting. JetSpec++ can also draft without pipelining when its acceptance-length gain outweighs the added draft time.

SELF-REFINEMENT EXAMPLE

  • Anchor replacement uses the first disagreement between the decoder’s argmax and the draft to write a replacement token and fix that position.The next pass reads refreshed states only through the accepted prefix and redrafts positions after the last fixed one.Source: D SELF-REFINEMENT EXAMPLE AND DRAFTING… · p. 24

DRAFTING WITHOUT PIPELINING

SELF-REFINEMENT EXAMPLE

  • In the illustrated refresh, the decoder agrees with the draft at positions 1–3, replacement fixes position 4, and the next pass independently fixes position 7.The figure depicts one batch row and follows Figure 3’s layout.Source: Figure 4 · p. 25

DRAFTING WITHOUT PIPELINING

  • Without pipelining, JetSpec++ is useful when its higher acceptance length over DSpark outweighs extra draft time, particularly when target verification takes most of the round.The comparison measures self-refinement without pipelining.Source: D SELF-REFINEMENT EXAMPLE AND DRAFTING… · pp. 24–25

E TRAINING AND EVALUATION DETAILS

The appendix specifies JetSpec++ training, model and baseline configurations, benchmark protocols, metrics, and correctness checks. Evaluation controls include a rule for handling repetition loops and a check against autoregressive outputs and eager draft execution.

Training data and optimization

  • The draft model trains for six epochs on 800k sequences drawn from Nemotron-Post-Training-Dataset-v2 and CodeAlpaca-20k, with AdamW and a peak learning rate of 3 × 10−4.Training takes 347 GPU-hours for Qwen3-8B and 710 for Qwen3.8-27B on GB200 GPUs; ablation drafts train two epochs on 239k mathematics sequences.Source: E TRAINING AND EVALUATION DETAILS · p. 26

Training blocks

  • Training blocks use up to 512 anchor positions per sequence and a full 16-position block; the clean prefix is fully noised with probability 0.25 or has each length 1–15 with probability 0.05.Each noise token copies a uniformly random token from the sequence up to and including the anchor.Source: E TRAINING AND EVALUATION DETAILS · p. 26

Draft model architecture and objective

  • The draft model has five transformer layers, uses target hidden states from five layers, attends causally within the block, and reuses the frozen target embedding and LM head.It has 1.05B trainable parameters for Qwen3-8B and 1.73B for Qwen3.8-27B; training uses forward KL at temperature 1 averaged over supervised draft positions without blockwise decay.Source: E TRAINING AND EVALUATION DETAILS · p. 26

Baselines

  • DFlash, DSpark, Domino, and JetSpec are trained on the same 800k sequences, with baseline-specific attention, loss, positional-decay, and batch settings.DFlash and JetSpec use five layers and the same five target layers, but fill every draft position with a mask token; DSpark uses 15 draft tokens after the anchor.Source: E TRAINING AND EVALUATION DETAILS · p. 26

Repetition loops and correctness checks

  • Outputs reaching 2,048 tokens are treated as repetition loops when their final 64-token span appears at least three times without overlap; runs stop at the loop’s onset.On GSM8K, 2,048 generated tokens from eight prompts match autoregressive output token by token in fp32, and 492 CUDA-graph draft rounds match eager recomputation.Source: E TRAINING AND EVALUATION DETAILS · p. 27

Benchmarks, metrics, and serving protocol

  • Acceptance length τ counts committed tokens per verification round, while throughput counts tokens per second across requests; hit-rate and hit/miss τ exclude each request’s first round.The evaluation fixes benchmark shuffling and generation settings, and Table 3 uses batch size 1 with speedups relative to the matching autoregressive row.Source: E TRAINING AND EVALUATION DETAILS · pp. 26–27

F THEORETICAL ANALYSIS

The appendix gives three theoretical results supporting the method: exact verification-outcome probabilities under a perfect decoder, a guarantee for per-round outcome selection, and lossless verification after anchor replacement.

Theoretical results

  • With a perfect decoder, the score P(k, a) equals the probability of verification outcome (k, a) at any temperature.The appendix states this result under its stated assumptions.Source: F THEORETICAL ANALYSIS · p. 27
  • With a perfect decoder, per-round selection by Eq. 7 is at least as good as any fixed allocation.The selection ranks outcomes by their decoder score.Source: F THEORETICAL ANALYSIS · p. 27
  • Verification remains lossless for every draft, including drafts rewritten by anchor replacement.The appendix identifies this as a result supporting Section 4.3.Source: F THEORETICAL ANALYSIS · p. 27

F.1 EXACT OUTCOME PROBABILITIES WITH A PERFECT DECODER

Under the perfect-decoder assumption, the decoder score gives the exact probability of each possible verification outcome. This result supports selecting cached outcomes by their predicted verification likelihood.

Outcome selection

  • A perfect decoder’s score P(k, a) is the exact verification-outcome probability, so Eq. 7 selects the cache with the largest hit probability.The stated result holds at any temperature and for a draft fixed before verification.Source: §F.1 · p. 27

Example

  • For γ = 2, the example gives P(1, b) = 0.24, matching the probability of accepting d1, rejecting d2, and drawing b from the residual.The example uses target probabilities 0.8 for d1, 0.5 for d2 after d1, and 0.3 for b after d1.Source: §F.1 · pp. 27–28

Assumptions

  • The result assumes a perfect decoder whose estimates equal the target next-token distribution, fixed draft tokens before verification, independent verification draws, and point-mass proposals.The draft variants described for Sections 4.1–4.3 satisfy the fixed-draft and draw-independence conditions.Source: §F.1 · p. 28

Exact outcome probabilities

  • Under Assumption 1, P(k, a) equals the probability of outcome (k, a), for any temperature and every fixed draft.For k < γ, the outcome requires a ≠ d_{k+1}; the γ-acceptance case uses P(γ, a) = P(γ) ˆp(a | γ).Source: §F.1 · p. 28

F.2 PER-ROUND SELECTION AND THE GEOMETRIC FAN-OUT OF SSD

The analysis compares adaptive per-round cache selection with SSD’s fixed fan-out and establishes lossless verification for argmax drafts using a point-mass proposal.

F.2 PER-ROUND SELECTION AND THE GEOMETRIC FAN-OUT OF SSD

  • Under SSD’s assumptions, the continuous surrogate recovers SSD’s geometric fan-out only on retained active lengths, which need not include every length.The proposition’s surrogate uses an acceptance-rate law and a power-law hit-rate curve.Source: §F.2 · pp. 28–29
  • With α=0.6, ρ=1, γ=4, and B=10, the all-length series captures 0.535 hit mass, versus 0.567 for the maximizing allocation omitting length 3.The maximizing allocation uses lengths 0, 1, 2, and 4.Source: §F.2 · p. 29
  • Per-round posterior selection is at least as good as any feasible cache choice in every round; with time-invariant posterior ranks, the selection is itself a fixed allocation.The best fixed allocation ties with that selection when the posterior ranks do not change across rounds.Source: §F.2 · pp. 29–30

F.3 LOSSLESS VERIFICATION

  • Every committed token is the target’s own draw at any temperature, whether the draft token came from the draft model or the decoder.For a fixed draft, verification uses a point-mass proposal at each draft token and accepts it with the target probability.Source: §F.3 · pp. 30–31
  • Using the draft model’s softmax as the proposal for an argmax token breaks the lossless guarantee because that token was not sampled from the softmax.The passage gives an example where the resulting commitment probability differs from the target probability.Source: §F.3 · p. 31

Related papers

Loading 2610.19792v2…