JetSpec++: Foresight Parallel Drafting with Self Refinements

Paper | Webpage | Models

Overview

JetSpec++ is a lossless speculative decoding method that hides almost all of its drafting cost behind verification. It drafts the next block on a second GPU while the target model verifies the current one, and it refines that draft while it waits.

• Faster: on a GB200 pair, up to 7.45x over autoregressive decoding and 1.54x over DSpark, at up to 1,361 tokens per second for one request.
• Hidden drafting: the time per round beyond verification falls by up to 84% relative to DSpark.
• Longer acceptance: the longest acceptance length on all seven benchmarks, 8.94 tokens per round on average at temperature 0, against 6.89 for DSpark.

jetspecpp-overview.mp4

Decode speed on a GB200 pair at batch size 1 and temperature 0, with the speedup over autoregressive decoding (AR):

Target	Benchmark	AR (tok/s)	DSpark (tok/s)	JetSpec++ (tok/s)
Qwen3-8B	GSM8K	182.3	744.3 (4.08x)	1148.8 (6.30x)
Qwen3-8B	MATH-500	182.7	902.1 (4.94x)	1361.1 (7.45x)
Qwen3.8-27B	GSM8K	60.3	249.6 (4.14x)	342.4 (5.68x)
Qwen3.8-27B	MATH-500	59.3	279.8 (4.72x)	367.2 (6.20x)

Supported Models

Target	Checkpoint	GPUs
Qwen3-8B	JetSpec/jetspec-plus-plus-qwen3-8b	A100, B200, GB200
Qwen3.8-27B	JetSpec/jetspec-plus-plus-qwen3p8-27b	A100, B200, GB200

Installation

JetSpec++ runs on Linux with two GPUs that have peer access and needs a CUDA compiler. Install the measured environment for your platform, then the package:

git clone https://github.com/hao-ai-lab/JetSpec.git && cd JetSpec
pip install --no-deps -r jetspec/engines/jetspec_plus_plus/requirements-x86_64.lock
pip install --no-deps -e .

On GB200, use requirements-gb200-aarch64.lock with --extra-index-url https://download.pytorch.org/whl/cu128.

Quick Start

python -m jetspec.engines.jetspec_plus_plus generate --model qwen3-8b \
    "How many positive whole-number divisors does 196 have?"

The command downloads the pinned target and checkpoint, picks the schedule for your GPU, and prints the answer with its throughput. That throughput covers one short prompt on a fresh engine, so use speedup.py below for a steady number. The first run also compiles kernels into ~/.cache/jetspec-plus-plus, which takes several minutes. --devices names two GPUs as nvidia-smi numbers them (default 0,1). examples/jetspec_plus_plus/ has a launch script for each model.

Evaluation

Speedup over autoregressive decoding on your GPU. --dataset takes gsm8k, math500, aime25, humaneval, mbpp, livecodebench or mt-bench:

python bench/jetspec_plus_plus/speedup.py --model qwen3-8b --dataset math500

The paper's Table 3 on a GB200 pair, against DSpark. The guide explains the host file:

python bench/jetspec_plus_plus/reproduce_paper_results.py --model qwen3-8b --assets release --host-values host-8b.json
python bench/jetspec_plus_plus/reproduce_paper_results.py --model qwen3.8-27b --assets release --host-values host-27b.json

Code

• Method: the torch reference holds the method in readable form.
• Engine: the serving engine runs the method on two GPUs.

JetSpec

Overview

Paper | Blog | Webpage | Models

JetSpec is an implementation of causal parallel tree drafting for fast LLM speculative decoding inference with up to 10x acceptance length, and 1000+ TPS on coding and math tasks using B200 GPUs. A causal-parallel draft head proposes a token tree, and the frozen target model verifies the whole tree in one forward pass under a tree-causal attention mask. The accepted path is selected in accordance with the target's own logits, so decoding is lossless by construction.

Supported Models

Target family	Draft head
Qwen3-8B	JetSpec/jetspec-qwen3-8b
Qwen3-30B-A3B	JetSpec/jetspec-qwen3-30b-a3b
Qwen3.6-35B-A3B	JetSpec/jetspec-Qwen3.6-35B-A3B
Gemma4-26B-A4B0-it	JetSpec/jetspec-gemma4-26B-A4B0-it
GPT-OSS-20B	JetSpec/jetspec-gpt-oss-20b
Step3p7-Flash	JetSpec/jetspec-Step-3.7-Flash

Most benchmark and diagnostic scripts can read the draft head from JETSPEC_DRAFT_HEAD:

export JETSPEC_DRAFT_HEAD=JetSpec/jetspec-qwen3-8b

Installation

pip install -e '.[bench,kernel]'

For FlashAttention 2 benchmarks, install the extra after build dependencies are available:

pip install -e '.[bench,flash-attn]'

If you use uv, the project includes extra build dependency metadata for flash-attn.

Quick Start

JetSpec has two execution paths: jetspec/core/, a lightweight HuggingFace-based reference implementation, and jetspec/engines/jetspec/, an optimized serving engine with paged KV, custom Triton tree attention, and CUDA graphs for better wall-clock latency and throughput.

HF reference:

from jetspec import LLM, SamplingParams

llm = LLM("Qwen/Qwen3-8B", attn_implementation="flash_attention_2")
out = llm.generate(
    "The three primary colors are",
    SamplingParams(temperature=0.0, max_new_tokens=64),
)
print(out["text"])

Inference engine, with paged KV, a Triton tree-attention backend, compiled verification, CUDA graphs, and optional session reuse:

from jetspec import load_draft_head, DraftHeadTreeDrafter
from jetspec.engines.jetspec import JetSpecEngine, SamplingParams

engine = JetSpecEngine(
    "Qwen/Qwen3-8B",
    attn_backend="triton_paged_tree_compiled",
)
head = load_draft_head("JetSpec/jetspec-qwen3-8b")
drafter = DraftHeadTreeDrafter(
    head,
    target=engine.model,
    block_size=head.block_size,
    target_layer_ids=head.target_layer_ids,
)
out = engine.generate_tree(
    "The three primary colors are",
    drafter,
    block_size=head.block_size,
    tree_depth=20,
    tree_width=7,
    budget=128,
    target_layer_ids=head.target_layer_ids,
    sampling_params=SamplingParams(temperature=0.0, max_new_tokens=64),
)
print(out["text"])
print("tokens per forward:", out["tpf"])

Evaluation

HF reference benchmark. With --include-dflash-baseline, it reports the shared AR-greedy baseline, the linear DFlash block baseline using the draft head's block_size, and the JetSpec tree construction algorithms selected through --algos:

CUDA_VISIBLE_DEVICES=0 \
JETSPEC_DRAFT_HEAD=JetSpec/jetspec-qwen3-8b \
python bench/jetspec/reference/benchmark.py \
  --model Qwen/Qwen3-8B \
  --attn-implementation flash_attention_2 \
  --tree-attn-implementation triton \
  --dataset gsm8k \
  --samples 64 \
  --algos accum_logp \
  --depth 20 \
  --width 7 \
  --budget 256 \
  --max-new 1024 \
  --warmup-rounds 3 \
  --include-dflash-baseline

Optimized engine, wall-clock latency and throughput. The headline throughput uses --drafter graphed, a CUDA-graph replay of the draft-head forward per context-capacity bucket. --drafter eager is correct but about 2x slower:

CUDA_VISIBLE_DEVICES=0 \
JETSPEC_FUSE_GEMMS=1 \
JETSPEC_BACKEND=triton_paged_tree_cudagraph_nogather \
JETSPEC_DRAFT_HEAD=JetSpec/jetspec-qwen3-8b \
python bench/jetspec/engine/tps_walltime.py \
  --prompt-set gsm8k \
  --samples 64 \
  --max-tokens 2048 \
  --tree-depth 20 \
  --budget 128 \
  --drafter graphed \
  --warm-all \
  --session

The vLLM v1 integration supports MATH-500 testing through <vllm-jetspec>/examples/offline_inference/jetspec_profiling_math500_tree_budget_bsz_sweep.sh and HumanEval testing through <vllm-jetspec>/examples/offline_inference/jetspec_profiling_humaneval_tree_unit_kvlayout.sh.

Engine results: single-stream (batch size 1) Qwen3-8B on NVIDIA B200, bf16, --drafter graphed, warm steady state, with a pinned CPU reservation.

dataset	JetSpec engine TPS	accept_len
MATH-500	1150 tok/s	9.56
GSM8K	984 tok/s	7.94
HumanEval	867 tok/s	6.92
AIME25	849 tok/s	8.79
MBPP	789 tok/s	7.61
LiveCodeBench	684 tok/s	7.66
MT-Bench	545 tok/s	4.88

Tree algorithms exposed by bench/jetspec/reference/benchmark.py:

Algorithm	Purpose
accum_logp	robust breadth-first cumulative-logprob tree; default baseline
top2gap_fanout	adaptive fanout using the per-depth top-2 logprob gap
task_router	prompt/task-aware routing over tree shapes
reasoning_router	reasoning-pattern-aware routing
class_histogram	class-conditioned profile-guided tree shaping
depth_rank_histogram	offline profile table over (depth, rank) acceptance

Code

• Method: jetspec/methods/jetspec/ holds tree drafting.
• Engine: the serving engine adds paged KV, Triton tree attention and compiled verification.

For Coding Agents

AGENTS.md describes every command, option and component in this repository. The method and engine folders each have their own AGENTS.md.

Citation

@article{feng2026jetspecpp,
  title   = {JetSpec++: Foresight Parallel Drafting with Self Refinements},
  author  = {Feng, Zhaoxiang and Hu, Lanxiang and Yuan, Haoran and Qian, Yu-Yang and Yang, Jihui and Zhao, Peng and Zhu, Yibo and Rosing, Tajana and Zhang, Hao},
  journal = {arXiv preprint arXiv:2610.19792},
  year    = {2026}
}

@article{hu2026jetspec,
  title   = {JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting},
  author  = {Hu, Lanxiang and Feng, Zhaoxiang and Wu, Yulun and Yuan, Haoran and Zhao, Yujie and Qian, Yu-Yang and Wang, Bojun and Zhao, Peng and Jiang, Daxin and Zhu, Yibo and Rosing, Tajana and Zhang, Hao},
  journal = {arXiv preprint arXiv:2606.18394},
  year    = {2026}
}