PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Lanxiang Hu, Zhaoxiang Feng, Yulun Wu, Haoran Yuan, Yujie Zhao, Yu-Yang Qian, Bojun Wang, Peng Zhao, Daxin Jiang, Yibo Zhu, Tajana Rosing, Hao Zhang

arXiv:2606.18394v3cs.CL

TL;DR

Speculative decoding struggles to scale because larger draft budgets can increase drafting overhead or reduce acceptance. JetSpec uses a causal parallel draft head to generate branch-consistent trees efficiently, consistently outperforming baselines and reaching 9.64× speedup on MATH-500.

  • Problem

    Speculative decoding faces a causality–efficiency dilemma: autoregressive drafting is costly at greater tree depths, while one-pass bidirectional drafting can produce mutually inconsistent trees and lower acceptance.

  • Method

    JetSpec trains a causal parallel draft head that generates multiple tree nodes in one pass while preserving branch-wise causal conditioning aligned with the target model’s autoregressive factorization.

  • Results

    Across dense and MoE Qwen3 models and serving settings, JetSpec consistently improves speedup over baselines, reaching 9.64× on MATH-500 on Qwen3-8B.

  • Takeaways & Limitations

    JetSpec shows that larger speculative-decoding budgets can translate into higher wall-clock speedup while reducing decoding latency across evaluated workloads.

  • Takeaways & Limitations

    JetSpec evaluates static budget policies and leaves dynamic serving-time budget scheduling to future work.

Abstract

from arXiv · show

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.

1 Introduction

Speculative decoding addresses serial LLM decoding with low adaptation cost, but head-based methods trade path-conditioned draft quality against one-pass efficiency. JetSpec resolves this dilemma through causal parallel tree drafting, scaling acceptance and speed with larger draft budgets while demonstrating strong benchmark and serving gains.

  • Motivation: Serial decoding creates latency challenges for long-generation applications including math, coding, and agentic reasoning.Autoregressive generation remains largely sequential, making latency a major deployment bottleneck.
  • Challenge: Head-based speculative decoding faces a causality-efficiency dilemma between sequential path-conditioned drafting and one-pass branch-agnostic drafting.Autoregressive drafters require sequential passes as tree depth grows, whereas bidirectional drafters can produce mutually inconsistent trees that waste budget and reduce acceptance.
  • JetSpec: JetSpec trains a causal parallel draft head that predicts multiple tree nodes in one pass while preserving branch-wise causal conditioning.Block-level causal attention over hidden states aligns each branch’s draft distribution with the target model’s autoregressive factorization.
  • Results: τ = 10.7 and 9.5× speedup on MATH-500 were achieved with 256 draft tokens, while MT-Bench reached τ = 5.9 and over 4× speedup.JetSpec was competitive with 16 draft tokens and achieved stronger gains in the high-token-budget regime.
  • Results: More than 9.5× end-to-end speedup was achieved by scaling JetSpec’s draft token budget.The method is designed to jointly optimize drafting cost and acceptance rate while leveraging additional decoding compute.
  • Serving evaluation: JetSpec consistently outperforms baselines under small to moderate serving loads after integration into an industry-grade serving engine.The serving evaluation uses realistic request-rate scenarios and demonstrates practical latency-reduction potential.

2 JETSPEC

JetSpec targets speculative decoding’s scaling bottleneck by jointly lowering drafting cost and improving acceptance through causal parallel tree drafting. It reuses frozen target-model features and applies tree-causal conditioning so candidate branches remain aligned with autoregressive target factorization while being generated in parallel.

  • Scaling Bottlenecks in Speculative Decoding: Longer drafts improve throughput only when acceptance rate α stays high and cumulative drafting overhead Nc remains small.As draft length N grows, modest acceptance degradation or non-negligible drafting cost can erase parallel-verification benefits.
  • Limitations of Existing Techniques: Existing methods typically improve either acceptance quality α through alignment or drafting cost c through head-based and block-parallel drafting.The paper argues that favorable scaling requires high α and low c simultaneously, motivating JetSpec’s joint approach.
  • Can We Enable Both?: Branch-agnostic parallel drafting can produce individually plausible but mutually inconsistent continuations, lowering tree acceptance despite lossless verification.Without conditioning on selected ancestor tokens, tree construction ranks branches under a pseudo-distribution rather than a causally valid draft distribution.
  • Architecture: JetSpec combines a head-based architecture with a causal parallel decoding head that reuses intermediate target-model features to generate candidate trees cheaply.Its design extracts fused hidden features from the frozen target model and injects target context into draft generation.
  • Architecture: The tree-causal attention mask lets all nodes run in parallel while restricting each node to the prefix and its own ancestors, mirroring autoregressive factorization.This branch-wise conditioning aligns the draft distribution with the target model while avoiding visibility of descendants or unrelated sibling branches.
  • Tree Construction and Verification: The candidate tree is built by repeatedly expanding the highest-scoring node with up to W children until the draft budget is filled or no expandable node remains.Verification then accepts the longest prefix whose draft tokens satisfy the target model’s acceptance condition.

3 Experiments

Experiments show that JetSpec consistently outperforms tree-based baselines, with larger draft budgets improving speedup and causal drafting remaining effective under non-greedy decoding. Serving and ablation results further show that performance depends on load, training choices, target-model generalization, and causal consistency.

  • Overall results: JetSpec consistently outperforms tree-based baselines across greedy and non-greedy decoding evaluations.The evaluation covers math, coding, and chat benchmarks on Qwen3 target models.
  • Draft-budget scaling: At high draft budgets, JetSpec improves from roughly 4–6× to 7–10× speedup on math and coding benchmarks, while DDTree scales only up to 9×.The results attribute JetSpec’s stronger scaling to causal tree drafting and report robustness under temperature 1.
  • Serving evaluation: At batch size 1, increasing the tree budget from 16 to 128 raises throughput from 443.3 to 968.2 TPS and speedup from 3.09× to 6.75×.At higher batch sizes, larger budgets saturate as verification cost and compute pressure reduce their benefit.
  • Training ablations: A learning rate of 3 × 10−4 reaches peak speedup of 8.30×, while reverse-KL distillation causes a 36–46% relative drop versus forward-KL.SFT and forward-KL remain within ∼3% across the four math benchmarks.
  • Generalization and data: JetSpec generalizes to the Qwen3-30B-A3B MoE target with competitively higher speedups than DDTree, while regenerated target-model sequences provide the strongest training performance.Corpus-trained JetSpec still achieves consistent speedups across benchmarks.
  • Causal versus diffusion heads: Across MATH-500 prompts, the diffusion head has a rank-1 gap exceeding the causal head’s on 92% of prompts, with a median gap 5× larger.The mean accepted-token ratio is 1.95×, closely matching the 1.52× macroscopic speedup ratio.

4 Related Work

Prior speculative decoding work improves speed through draft-target alignment, specialized head-based drafters, and draft-free retrieval. Related parallel decoding methods reduce autoregressive latency through early-layer self-speculation, diffusion refinement, or causal fixed-point updates.

  • Speculative Decoding: Speculative decoding drafts tokens for parallel target-model verification, with speed determined by draft quality and drafting cost.Prior approaches address these factors through draft-target alignment, head-based drafters, or draft-free retrieval methods.
  • Drafter Alignment: Drafter-alignment methods train draft models or heads using future-token prediction, fine-tuning, distillation, or acceptance-rate optimization.These methods primarily improve agreement between draft and target distributions, thereby targeting the alignment term controlling acceptance.
  • Parallel Decoding and Drafting: Parallel decoding reduces latency by relaxing one-token-at-a-time generation through self-speculation, diffusion refinement, or causal fixed-point updates.Consistency distillation has been used to improve convergence of Jacobistyle methods toward autoregressive outputs.

5 Conclusion

JetSpec is a causal parallel tree-drafting framework that jointly targets drafting latency and accepted length. It preserves branch-wise causal dependencies while generating candidate trees in one low-cost draft-head pass aligned with the target model’s autoregressive factorization.

  • 5 Conclusion: JetSpec jointly addresses drafting latency and accepted length, the two factors limiting speculative-decoding speedup.The framework is presented as a causal parallel tree-drafting approach.
  • 5 Conclusion: JetSpec preserves branch-wise causal dependencies while generating candidate trees in a single low-cost draft-head pass.This design combines causal structure with efficient tree construction.
  • 5 Conclusion: JetSpec makes tree construction better aligned with the target model’s autoregressive factorization.The alignment follows from preserving branch-wise causal dependencies during parallel drafting.
  • 5 Conclusion: The framework is evaluated across dense and MoE Qwen3 models and vLLM serving settings.The supplied conclusion passage identifies these model and serving contexts without reporting their quantitative outcomes.

A Tree Drafting Quality Analysis · A.1 Setup

The analysis compares causal and diffusion heads by measuring draft-surrogate and target conditional probabilities along constructed trees, showing that branch-wise causal conditioning preserves coherent high-ranked paths and improves verification acceptance. Experiments use matched heads and a fixed tree-construction and verification setup on MATH-500.

  • A Tree Drafting Quality Analysis: The causal head’s faithful rank-1 branch is accepted for 6 tokens, whereas the diffusion head’s incoherent rank-1 branch is accepted for only 4 tokens.The causal branch “ are told that” aligns with the target model’s joint log probability, while the diffusion branch composes mutually inconsistent tokens.
  • A Tree Drafting Quality Analysis: The diffusion head’s rank-1 branch has target joint log probability −63.32 nats, corresponding to probability approximately e−63.Its branch-agnostic per-position predictions compose “ given” at depth 1 with “ told” at a later depth, producing an incoherent path.
  • A Tree Drafting Quality Analysis: The case study uses MATH-500 prompt #0 at decode step 0, with both heads drafting from the same prefix whose last token is “We”.Figure 4 visualizes the rank-1 versus rank-3 contrast, while the appendix provides full top-5 branches and rank-1 gap statistics.
  • A.1 Setup: For each tree node v, the analysis records cumulative draft-surrogate and target conditional log probabilities together with the verifier’s accepted path.These measurements compare draft scores with target probabilities conditioned on the node’s preceding path.
  • A.1 Setup: The evaluation uses FlashAttention-2 and Triton tree-attention within a modified speculative-decoding loop that dumps the constructed tree at a chosen decode step.This setup enables direct comparison between draft rankings, target conditionals, and verified paths.
  • A.1 Setup: Both heads share the prefill and produce N parallel position-wise distributions r1, . . . , rN for tree construction.The comparison therefore isolates differences between causal and diffusion drafting under a shared target-model prefix.
  • A.1 Setup: The experiments train both heads with learning rate 3×10−4 and γ = 0, evaluating them on 4×B200 with block size N = 16.The maximum draft depth is also N = 16.
  • A.1 Setup: Tree construction uses cumulative-log-probability ranking with branching width W = 7 and a node budget B.The passage specifies the branching width and introduces the node budget, but its displayed value is truncated.

A.2 Top-5 Branches at the Canonical Example … B Tree-Drafting Algorithm Details

The appendix shows that causal branch conditioning produces more coherent top-ranked trees and longer accepted prefixes than diffusion marginals, then formalizes JetSpec’s priority-queue tree-drafting procedure under a fixed node budget.

  • A.2 Top-5 Branches at the Canonical Example: At MATH-500 prompt #0, the diffusion tree’s top-2 and top-4 branches combine marginally plausible tokens with target joints below −50 nats.Only diffusion rank-3 is reported as coherent, with target joint −0.08.
  • A.3 Gap Distribution over 50 Prompts: 26% of diffusion prompts have rank-1 gaps ≥+80 nats versus 0% for causal, while faithful gaps below +5 nats occur in 6% versus 42%.The diffusion upper tail corresponds to joint target probability below 10−35.
  • A.3 Gap Distribution over 50 Prompts: At γ = 0, mean accepted length is 4.84 tokens for diffusion versus 9.46 for causal.
  • A.3 Gap Distribution over 50 Prompts: At γ = 7, diffusion tree quality recovers: its extreme-tail rate falls to 4% and faithful rank-1 rate rises to 26%.γ = 7 is identified as DFlash’s best macroscopic loss-weighting setting.
  • A.4 Why the Failure Mode Is Structural: Diffusion heads generate all position-wise distributions from one shared hidden state without causal masking, so later marginals do not condition on earlier tokens.This can promote syntactically incompatible token combinations into the top-ranked branch.
  • B Tree-Drafting Algorithm Details: Algorithm 1 constructs a candidate tree by repeatedly expanding the highest-scoring queued node, adding up to W children until depth N or node budget B is reached.Each child receives an updated path score and returns to the priority queue.

C Tree-Drafting Algorithm Ablation

The ablation shows that cumulative draft log-probability is essential for effective tree construction, while entropy-only scoring collapses performance. Low-entropy-weight hybrid scoring nearly matches the default, but larger weights degrade speedup and acceptance.

  • Cumulative log-probability scoring: 8.15× speedup with τ = 9.81 makes cumulative draft log-probability the production default.The method expands branches best-first by cumulative draft log-probability, prioritizing high-likelihood continuations within the budget.
  • Entropy-only collapse: 4.76× speedup represents a −42% relative drop from accum_logp when scoring uses per-depth entropy alone.The collapse confirms that marginal entropy alone cannot reliably identify high-acceptance branches.
  • Hybrid scoring: At α ≤1, hybrid scoring stays within ∼1.5% of accum_logp at 8.15–8.27× speedup.As α increases, acceptance degrades monotonically, reaching 7.42× at α=8.
  • Hybrid scoring: The α →0 hybrid limit reduces to cumulative log-probability, whereas α →∞ approaches the 4.76× entropy-only collapse.These limits explain why low-α hybrid and accum_logp share a performance plateau while high entropy weighting is harmful.

D Training Details

JetSpec trains a causal parallel draft head on fused hidden states from a frozen target model, using block-local causal attention to predict masked future tokens in parallel. Training uses matched block size and standardized optimization settings for comparison with DFlash.

  • Optimization: A learning-rate sweep over 1×10−4 to 1×10−3 found 3×10−4 and 6×10−4 generally best, with 3×10−4 used by default.All draft-head training runs use 8 H100 GPUs with a micro batch of 2.
  • Draft-head training: JetSpec uses block size 16, retaining each block’s first token as an anchor while masking the remaining positions for parallel prediction.The anchor is retained as block context and excluded from the loss; loss applies only to future token positions within each block.
  • Draft-head training: Each query attends to the verified prefix and earlier positions in its own block, but not future positions or other sampled blocks.This causal mask preserves branch-wise conditioning while allowing all masked future tokens to be predicted in parallel.
  • Model implementation: For Qwen3-8B, JetSpec concatenates hidden states from target layers {1, 9, 17, 25, 33} and projects the resulting 5d feature to d = 4096.The target model remains frozen, while the fused features provide context to the lightweight causal draft head.
  • Model implementation: The draft head is a lightweight Qwen3-style decoder with 5 layers, 32 attention heads, and 8 KV heads.The supplied implementation description also specifies the head dimension, but the passage is truncated before its value.

E vLLM Implementation Details

JetSpec integrates causal parallel tree drafting into vLLM, using causal attention to generate tree-organized candidates from target hidden states. vLLM verifies all speculative nodes in one tree-attention pass, while tree-budget benefits depend on serving batch size.

  • Drafting: JetSpec implements causal parallel tree drafting inside vLLM, with causal attention and candidates organized as a speculative tree rather than a linear block.This combines one-pass drafting with tree-structured candidate organization.
  • Drafting: The proposer extracts top-k tokens and log-probabilities at each future position, then constructs a speculative tree under a fixed node budget.Supported expansion scores include cumulative log-probability, entropy-guided scoring, and hybrid scoring; experiments use cumulative log-probability.
  • Verification: The target verifies all speculative nodes in one forward pass using parent indices, depths, and a tree attention mask.An optimized fused paged tree-attention kernel applies the mask without materializing a dense per-request mask.
  • Serving results: Larger tree budgets improve low-batch latency by reducing verification rounds, but their relative gain decreases at larger batch sizes as verification and memory pressure increase.Table 11 reports this serving trade-off on HumanEval with a single H100 GPU.

F Ablation Study Details

The ablations show that JetSpec is robust to key training choices and generalizes beyond dense Qwen3-8B. Performance is best supported by moderate learning rates, distribution-preserving objectives, regenerated target-model data, and causal parallel drafting across architectures.

  • Experimental setup: Ablations use Qwen3-8B on the Nemotron Post-Training Dataset V2 math split with block size 16 and tree node budget 255.Tree construction uses accumulated draft log-probability and the main-experiment evaluation protocol.
  • Learning Rate: Performance plateaus around 3 × 10^-4 learning rate, while very small rates underfit and larger rates add no consistent gains.The drafter is underfit at very small learning rates; larger rates remain competitive without reliable improvement.
  • Training Objective: SFT and forward-KL perform similarly, whereas reverse-KL substantially degrades budgeted tree-drafting performance.The result suggests mode-seeking distillation poorly preserves the multiple plausible continuations needed for tree drafting.
  • Model Generalizability: On Qwen3-30B-A3B MoE, JetSpec continues to outperform DDTree under the same training recipe.This shows causal parallel tree drafting is not specific to a single model architecture.
  • Training Data: Regenerated target-model sequences produce the strongest performance, while corpus-trained JetSpec is weaker but still effective.Regenerated data better matches the target model’s own generation distribution.

G Empirical Per-token Drafting Cost on Modern Hardware

On a single H200 NVL GPU, per-draft-token cost falls as draft depth increases because parallel draft-head latency is amortized across proposed tokens. In the practical regime, this cost is below 1% of one target verification pass, although long-draft scaling also requires sufficiently high acceptance.

  • Results: Per-draft-token cost stays below 1% of one target verification pass for L ≤2048 and N ≥16, decreasing roughly as 1/N.At L=1024, c drops from 0.845% at N =16 to 0.054% at N =256.
  • Results: At L=1024, c falls from 0.845% at N =16 to 0.054% at N =256.The decrease reflects amortization of one parallel draft-head forward pass across N proposed tokens.
  • Results: Longer contexts increase c mildly, but it remains below 0.1% for N ≥256.The measurements sweep context lengths L ∈{128, . . . , 4096} and draft depths N ∈{1, 2, 4, 8, 16, 32, 64, 128, 256, 512}.
  • Implication: Across the same hardware, c decreases from ≈6.7% at N =2 to ≈0.05% at N =256.This moves drafting from the typical-cost regime toward the ultra-low-cost regime in Fig. 2.
  • Implication: Low drafting cost enables long-draft scaling only when acceptance remains sufficiently high.This condition motivates causal tree drafting as a way to improve scaling beyond cost reduction alone.
Loading 2606.18394v3…