Speculative decoding has become one of the more practically useful techniques for reducing LLM inference latency without touching model weights or degrading output quality. The core idea is straightforward: a cheap draft model proposes multiple tokens, and the target model verifies them in a single parallel pass, accepting a prefix of the proposal. The challenge is that this speedup saturates quickly. Increasing the draft budget only helps if acceptance rates stay high and drafting overhead stays low, and prior methods have struggled to achieve both simultaneously. JetFlow is a new framework that directly targets this scaling ceiling.
The Causality-Efficiency Dilemma
To understand what JetFlow is solving, it helps to be precise about why existing head-based methods plateau. Autoregressive drafters like EAGLE generate path-conditioned candidates, meaning each draft token is conditioned on its ancestors in the tree. This produces high-quality, coherent branches, but the drafting cost grows with tree depth because each level requires a sequential pass. Bidirectional block-diffusion drafters like DFlash go to the opposite extreme: they generate all draft positions in one forward pass, keeping per-token drafting cost extremely low. The problem is that their marginal distributions are branch-agnostic. A token at depth three is predicted without explicit conditioning on what was proposed at depths one and two along that branch, so the resulting trees can contain individually plausible tokens that are mutually inconsistent. This wastes budget and suppresses acceptance length.
This is not a minor engineering issue. The mathematical consequence is that the draft distribution does not factorize the same way the target model does. When the target model verifies a branch, it applies its autoregressive factorization; a draft distribution that ignores within-branch dependencies will systematically disagree with the target at deeper nodes, capping the accepted length regardless of how many tokens are proposed.
What JetFlow Does Differently
JetFlow's central contribution is a causal parallel draft head that generates all tree nodes in a single forward pass while preserving branch-wise causal conditioning. The mechanism is a tree-causal attention mask applied within the draft head: each node can attend to the original prefix and its ancestors in the tree, but not to siblings or descendants. This is implemented as a block-level masked attention over hidden states fused from the frozen target model.
The result is a draft distribution that factorizes branch-by-branch in a way that mirrors the target model's autoregressive factorization. Concretely, for a path through the tree, the draft probability of that path is a product of conditional distributions, each conditioned on the concrete ancestor tokens along that specific branch. This is structurally different from block-diffusion drafting, where the conditioning is on the shared prefix only, not on the path-specific ancestors. Crucially, JetFlow achieves this without sequential draft passes: all tree nodes are computed in parallel in one head forward pass, so the per-token drafting cost remains in the same ultra-low regime as DFlash.
The draft head is trained on top of frozen target model hidden states using supervised fine-tuning on approximately 780K examples drawn from the Nemotron Post-Training Dataset V2, covering math, coding, and chat domains. The training objective aligns the draft head's output logits with the target model's next-token predictions along sampled continuations.
Empirical Results
The paper evaluates on Qwen3-8B (dense) and Qwen3-30B-A3B (MoE), across GSM8K, MATH-500, AIME25, HumanEval, MBPP, LiveCodeBench, and MT-Bench. The headline numbers are:
- Up to 9.64x end-to-end speedup on MATH-500 with Qwen3-8B on H100, using a 256-token draft budget
- Average accepted length of 10.7 tokens on MATH-500, meaning the target model is effectively accepting over ten tokens per verification pass on math problems
- 4.58x speedup on MT-Bench, a more conversational workload where acceptance rates are naturally lower
- Consistent outperformance over DFlash-T (the tree-based DFlash variant) across all benchmarks on both model sizes
- On Qwen3-30B-A3B, JetFlow achieves 9.45x on MATH-500 versus DFlash-T's 8.61x, showing the approach generalises to MoE architectures
The drafting cost analysis is particularly informative. At a draft depth of 256 and context length up to 2048, the per-draft-token cost is below 0.1% of a single target verification pass, decreasing roughly as 1/N as the parallel draft cost is amortised across more proposed tokens. This is what makes scaling the draft budget worthwhile: the overhead of proposing 256 tokens instead of 16 is negligible relative to the verification cost, so any improvement in accepted length translates almost directly into wall-clock speedup.
The paper also integrates JetFlow into vLLM and evaluates under realistic serving loads across request rates. JetFlow consistently outperforms baselines under low-to-moderate load, with particularly strong gains on B200 hardware. This is an important practical validation, since serving scenarios introduce batching dynamics that can erode the benefits of speculative decoding at high request rates.
Limitations and Open Questions
A few things are worth scrutinising. First, the gains are most pronounced on math benchmarks, where outputs tend to be long, structured, and repetitive enough that a well-trained draft head can achieve very high acceptance rates. The MT-Bench results, while still strong, show a more modest ~4.5x speedup, which is more representative of what general-purpose chat deployments might see. The gap between math and chat performance is not surprising, but it does suggest that the practical benefit is domain-dependent.
Second, JetFlow's advantage over DFlash-T, while consistent, is not always dramatic. On MT-Bench with Qwen3-30B-A3B, the difference is 4.33x versus 4.26x. The causal conditioning matters more when the draft budget is large and when task structure rewards coherent multi-token continuations, which is precisely the math and coding regime. Whether the training overhead of learning the causal head versus a simpler diffusion head is justified in lower-acceptance domains is a reasonable question.
Third, the paper focuses on non-thinking mode for Qwen3. Thinking mode, which involves chain-of-thought reasoning with potentially very long outputs, is arguably the highest-value target for inference acceleration. Whether JetFlow's acceptance rates hold under thinking-mode distributions, which can be more varied and exploratory, is not addressed here.
The broader implication is that the causality-efficiency tradeoff in speculative decoding is not fundamental: it is a design choice. JetFlow demonstrates that a single forward pass through a masked attention head can recover much of the acceptance quality of sequential autoregressive drafting. This opens space for further work on tree construction heuristics, training objectives that more directly optimise accepted length, and integration with multi-token prediction heads that are increasingly being included in base model pre-training.
Code and models are available at github.com/hao-ai-lab/JetFlow. The full paper is on arXiv.