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

2026-06-16

AI Summary

This work addresses the trade-off between causality and efficiency in conventional speculative decoding, which struggles to simultaneously achieve high acceptance rates and low computational overhead when scaling draft budgets. The authors propose JetFlow, a head-based speculative decoding framework that introduces, for the first time, a causal parallel draft head incorporating hidden states from the target model. By generating candidate trees in parallel that satisfy branching causality constraints, JetFlow significantly increases accepted token lengths while maintaining efficient forward computation. The approach integrates frozen target-model hidden-state training with tree-structured speculative decoding and is implemented within the vLLM inference engine. Experiments demonstrate that JetFlow substantially outperforms existing methods on both dense and MoE variants of Qwen3, achieving speedups of 9.64× on the MATH-500 benchmark and 4.58× on conversational tasks when deployed on H100 GPUs.

Abstract

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 JetFlow, a head-based SD framework that combines one-forward drafting efficiency with branch-wise causal conditioning. JetFlow 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 JetFlow 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, JetFlow consistently outperforms bidirectional-head and tree-based SD baselines. On H100 GPUs, JetFlow 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/JetFlow.

Problem

• speculative decoding
• scaling limitation
• causality-efficiency dilemma
• drafting overhead
• acceptance rate

Innovation

• Speculative Decoding
• Parallel Tree Drafting
• Causal Conditioning
• Draft Head
• LLM Acceleration