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

# JetFlow: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting
Source: [https://arxiv.org/html/2606.18394](https://arxiv.org/html/2606.18394)
Lanxiang Hu1Zhaoxiang Feng1Yulun Wu2Haoran Yuan3Yujie Zhao1Yu\-Yang Qian4Bojun Wang5Daxin Jiang5Yibo Zhu5Tajana Rosing1Hao Zhang1 1UC San Diego2Zhejiang University3UIUC4Nanjing University5StepFun

###### 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 proposeJetFlow, a head\-based SD framework that combines one\-forward drafting efficiency with branch\-wise causal conditioning\.JetFlowtrains 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 enablesJetFlowto 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,JetFlowconsistently outperforms bidirectional\-head and tree\-based SD baselines\. On H100 GPUs,JetFlowachieves up to9\.64×9\.64\\timesspeedup on MATH\-500 and4\.58×4\.58\\timeson 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](https://github.com/hao-ai-lab/JetFlow)\.

## 1Introduction

Modern autoregressive \(AR\) Large Language Models \(LLMs\) operate under a simple yet costly bottleneck: decoding remains largely serial\. This sequential generation process makes latency a major challenge when they are deployed in applications such as math\(Weiet al\.,[2022](https://arxiv.org/html/2606.18394#bib.bib132); Zhuet al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib133)\), coding\(Deng and others,[2025](https://arxiv.org/html/2606.18394#bib.bib128); Jainet al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib139)\)and other agentic reasoning tasks\(Patilet al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib136); Yaoet al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib137); Luet al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib138)\), where models often perform long generations\.

Existing approaches address this challenge by either incorporating multi\-token prediction \(MTP\) during pre\-training or mid\-training\(Gloeckleet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib140); DeepSeek\-AI,[2024](https://arxiv.org/html/2606.18394#bib.bib149),[2026](https://arxiv.org/html/2606.18394#bib.bib150); Huanget al\.,[2026](https://arxiv.org/html/2606.18394#bib.bib151); Xiaoet al\.,[2026](https://arxiv.org/html/2606.18394#bib.bib152)\), or post\-training methods that adapt pretrained LLMs into diffusion LLMs \(dLLMs\) at the costs of varying degrees of performance degradation\(Zhaoet al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib147); Wuet al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib148)\)\. In contrast, applying speculative decoding \(SD\) at inference time to reduce generation latency while preserving quality\(Leviathanet al\.,[2023](https://arxiv.org/html/2606.18394#bib.bib141); Chenet al\.,[2023](https://arxiv.org/html/2606.18394#bib.bib156)\)\.

Among these techniques, SD stands out for its low adaptation cost and simplicity\. On one hand, prior work shows that SD performance can be substantially improved through lightweight draft\-model alignment, oftentimes less than 1B training tokens\(Liuet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib164); Zhouet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib165); Goelet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib166)\)\. On the other hand, head\-based methods such as Medusa and EAGLE reduce deployment complexity by avoiding the conventional separate draft\-target model setup\(Caiet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib142); Liet al\.,[2024b](https://arxiv.org/html/2606.18394#bib.bib143)\)\. More recent head\-based methods such as EAGLE\-3\(Liet al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib144)\)and DFlash\(Chenet al\.,[2026](https://arxiv.org/html/2606.18394#bib.bib145)\)continue to improve drafting quality and reducing drafting cost for better end\-to\-end acceleration\. Despite these advances, head\-based SD still faces a causality\-efficiency dilemma\. Autoregressive drafters such as EAGLE produce high\-quality path\-conditioned candidates but require sequential draft passes as tree depth grows\. Bidirectional block\-diffusion drafters such as DFlash produces drafts in one pass, but their branch\-agnostic marginals can form individually plausible yet mutually inconsistent trees, wasting budget and reducing acceptance\.

We proposeJetFlow, a SD paradigm designed to break this causality–efficiency dilemma:JetFlowtrains a causal parallel draft head to predict multiple tree nodes in one forward pass while preserving branch\-wise causal conditioning through block\-level causal attention over hidden states\. This gives each branch a path\-conditioned draft distribution aligned with the target model’s autoregressive factorization, enabling tree acceptance to better scale with additional draft budget\.

Empirically,JetFlowdemonstrates significant efficiency gains\. In the low\-token\-budget regime with 16 draft tokens,JetFlowachieves competitive end\-to\-end speedup\. In the high\-token\-budget regime with 256 draft tokens,JetFlowachieves strong gains, reachingτ=10\.7\\tau=10\.7and9\.5×9\.5\\timesspeedup on MATH\-500, and remaining effective on MT\-Bench withτ=5\.9\\tau=5\.9and over4×4\\timesspeedup\. We further integrateJetFlowinto industry\-grade serving engine and evaluate it across request rates, where it consistently outperforms baselines under small to moderate serving loads, with especially strong gains on higher\-end GPUs such as B200\. These results highlight the practical potential ofJetFlowfor reducing decoding latency in serving\.

![Refer to caption](https://arxiv.org/html/2606.18394v1/x1.png)Figure 1:End\-to\-end decoding speedup over standard autoregressive decoding on H100 GPUs across math, coding, and chat benchmarks\. DFlash denotes the original block\-parallel drafting method, DDTree is tree\-based variant of DFlash, andJetFlowdenotes our method\. Both employ a tree budget of 256 tokens using Algorithm[1](https://arxiv.org/html/2606.18394#alg1)\.In a nutshell, our work makes the following contributions:

- •We proposeJetFlow, a paradigm designed to jointly optimize drafting cost and acceptance rate for speculative decoding\.JetFlowtrains a causal prediction head with parallel tree\-drafting capability, enabling low\-cost draft generation while preserving path\-conditioned dependencies among draft tokens\.
- •We develop and evaluate tree\-drafting algorithms that allow speculative decoding to better leverage additional decoding compute\. By scaling the draft token budget,JetFlowachieves up toτ=10\.7\\tau=10\.7and more than9\.5×9\.5\\timesend\-to\-end speedup\.
- •We integrateJetFlowinto an industry\-grade serving engine and evaluate it under realistic serving scenarios\.JetFlowconsistently outperforms baselines under small to moderate serving loads, highlighting its practical potential for reducing latency in real\-world deployments\.

## 2JetFlow

### 2\.1Background

##### Scaling Bottlenecks in Speculative Decoding\.

In speculative decoding \(SD\), a lightweight approximation modelMqM\_\{q\}first proposesNNdraft tokens, which are then verified in parallel by the target modelMpM\_\{p\}\(Leviathanet al\.,[2023](https://arxiv.org/html/2606.18394#bib.bib141)\)\. The target model accepts the longest prefix that is consistent with its own predictions and resumes the next iteration from the first rejected token\. The practical throughput of SD depends on both the acceptance rate and the relative drafting cost\. Formally, letα\\alphadenote the average acceptance rate and letccdenote the cost coefficient between the time for a single step ofMqM\_\{q\}and that ofMpM\_\{p\}\. Under the standard i\.i\.d\. acceptance assumption, the acceptance length per speculative iteration is

𝔼​\[\#​tokens\]=1−αN\+11−α,\\mathbb\{E\}\[\\\#\\mathrm\{tokens\}\]=\\frac\{1\-\\alpha^\{N\+1\}\}\{1\-\\alpha\},\(1\)whereNNis the number of draft tokens\. Accordingly, the expected walltime improvement over standard autoregressive decoding is

Speedup=1−αN\+1\(1−α\)​\(N​c\+1\)\.\\mathrm\{Speedup\}=\\frac\{1\-\\alpha^\{N\+1\}\}\{\(1\-\\alpha\)\(Nc\+1\)\}\.\(2\)This expression exposes a key scaling bottleneck in speculative decoding: increasing the draft lengthNNimproves throughput only when the acceptance rateα\\alpharemains high and the cumulative drafting overheadN​cNcstays small\. AsNNgrows, even modest degradation inα\\alphaor non\-negligible draft cost can quickly offset the benefit of parallel verification\. Figure[2](https://arxiv.org/html/2606.18394#S2.F2)visualizes that jointly reducing drafting cost and improving acceptance quality is essential for unlocking favorable scaling with longer drafts\.

![Refer to caption](https://arxiv.org/html/2606.18394v1/x2.png)\(a\)Typical SD \(c=0\.05c=0\.05\)\.
![Refer to caption](https://arxiv.org/html/2606.18394v1/x3.png)\(b\)Ultra low\-cost SD \(c=0\.0005c=0\.0005\)\.

Figure 2:Expected speculative decoding speedup scales as a function of drafting lengthγ\\gamma, under different per\-token drafting costsccand acceptance ratesα\\alpha\. Comparing the two panels shows that reducingccsubstantially improves the scalability of speculative decoding with respect toγ\\gamma, and increasingα\\alphafurther amplifies this effect\. The results highlight that pushing per\-token drafting costcclow and acceptance rateα\\alphahigh is critical for unlocking stronger scaling with longer draft lengths\.
##### Limitations of Existing Techniques\.

Existing work mainly focuses on improving these termsseparately\. To increase acceptance qualityα\\alpha, draft\-model alignment methods improve agreement betweenMqM\_\{q\}andMpM\_\{p\}\(Zhouet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib165); Liuet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib164); Goelet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib166)\)\. Orthogonally, tree\-based SD methods such as SpecInfer, Medusa, and EAGLE improve verification efficiency by expanding and verifying multiple candidate branches in parallel\(Miaoet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib169); Caiet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib142); Liet al\.,[2024a](https://arxiv.org/html/2606.18394#bib.bib158),[2025](https://arxiv.org/html/2606.18394#bib.bib144)\)\.

To reduce drafting costcc, head\-based methods avoid a separately deployed draft model\(Caiet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib142); Liet al\.,[2024b](https://arxiv.org/html/2606.18394#bib.bib143),[2025](https://arxiv.org/html/2606.18394#bib.bib144)\)\. More recently, DFlash pushes drafting overhead to an even lower regime by using block\-diffusion\-style drafting to generate multiple draft tokens in one draft pass and therefore very low per\-token drafting cost\(Chenet al\.,[2026](https://arxiv.org/html/2606.18394#bib.bib145)\), see Appendix[G](https://arxiv.org/html/2606.18394#A7)\. However, longer\-draft scaling requires both highα\\alphaand lowccat the same time, motivatingJetFlowthat jointly improves drafting efficiency and acceptance quality\.

### 2\.2Head\-based Causal Parallel Drafting

##### Can We Enable Both?

A natural direction to co\-optimize both aspects is to combine tree drafting for higher acceptance with parallel drafting for lower cost\. However, existing parallel drafting technique like DFlash can break the causal structure needed for effective tree construction: independently predicted future positions are not conditioned on the tokens selected along each branch, leading to inconsistent continuations\.JetFlowaddresses this challenge by training a causal parallel draft head that preserves branch\-wise dependencies while retaining the efficiency of parallel prediction\.

##### Why Causality Matters\.

A key requirement for parallel tree drafting is that each node distribution should be conditioned on its own branch prefix\. Consider a draft tree rooted at the original prefixxx\. Each nodevvcorresponds to a candidate tokenyvy\_\{v\}\. Letπ​\(v\)=\(yv1,…,yv\)\\pi\(v\)=\(y\_\{v\_\{1\}\},\\ldots,y\_\{v\}\)denote the draft\-token path from the root tovv, and letπ<v\\pi\_\{<v\}denote the ancestor tokens before nodevv\. Following standard speculative decoding notation, letppdenote the target\-model distribution induced by target modelMpM\_\{p\}, and letqqdenote the draft distribution induced by drafterMqM\_\{q\}\. Thenq\(⋅∣x,π<v\)q\(\\cdot\\mid x,\\pi\_\{<v\}\)denotes the distribution over the candidate token at nodevv, conditioned on the concatenated context\(x,π<v\)\(x,\\pi\_\{<v\}\)\.

Thus, each nodevvat depthiishould be sampled from the continuation\-aware distributionq​\(yv∣x,π<v\)q\(y\_\{v\}\\mid x,\\pi\_\{<v\}\), rather than from a branch\-agnostic per\-position distributionri\(⋅∣x\)r\_\{i\}\(\\cdot\\mid x\)that is not conditioned on the concrete ancestor tokens selected along the branch\. Without this causal conditioning, for a candidate branchy1:ky\_\{1:k\}, tree construction ranks branches according to a pseudo\-distribution

qsur​\(y1:k∣x\)∝∏i=1kri​\(yi∣x\),q\_\{\\mathrm\{sur\}\}\(y\_\{1:k\}\\mid x\)\\propto\\prod\_\{i=1\}^\{k\}r\_\{i\}\(y\_\{i\}\\mid x\),\(3\)whererir\_\{i\}denotes the branch\-agnostic draft distribution at positionii\. Thus, tree construction can favor continuations whose tokens are individually plausible but mutually inconsistent\. Verification remains lossless by rejecting invalid prefixes, but the tree itself is optimized underqsurq\_\{\\mathrm\{sur\}\}rather than a causally valid draft distribution, leading to lower acceptance\. We provide an empirical case study in Section[3\.4\.2](https://arxiv.org/html/2606.18394#S3.SS4.SSS2), with qualitative example in Appendix[A](https://arxiv.org/html/2606.18394#A1)\.

##### Architecture\.

JetFlowadopts a head\-based drafting architecture: the draft head can reuse rich intermediate features from the target model while keeping draft generation cheap\. In particular, DFlash shows that a lightweight block\-parallel draft head can reuse target\-model hidden states and inject fused target\-context features into the draft layers’ KV cache for strong target\-model guidance\(Chenet al\.,[2026](https://arxiv.org/html/2606.18394#bib.bib145)\)\. Building on this design,JetFlowintroduces a causal parallel decoding head for tree drafting\.

![Refer to caption](https://arxiv.org/html/2606.18394v1/x4.png)Figure 3:JetFlowdesign overview\.JetFlowextracts fused hidden features from the frozen target model and conditions a causal\-parallel draft head to generate high\-quality candidate trees in one forward pass\.Given a prefixxx, with a corresponding fused hidden state featureshxoh\_\{x\}^\{o\}, the draft head predicts logits for multiple active tree nodes in parallel\. Following the standard speculative decoding notation, letMpM\_\{p\}denote the target model andMqM\_\{q\}denote the draft model or draft head, with corresponding distributionsppandqq\. The target model defines an autoregressive distribution over a length\-kkcontinuation:

p​\(y1:k∣x\)=∏i=1kp​\(yi∣x,y<i\)\.p\(y\_\{1:k\}\\mid x\)=\\prod\_\{i=1\}^\{k\}p\(y\_\{i\}\\mid x,y\_\{<i\}\)\.\(4\)To align parallel drafting with this causal factorization, we apply a tree\-causal attention mask to the draft head\. Each node can attend to the original prefix and its ancestors, but not to descendants or unrelated sibling branches\. For two tree nodesuuandvv, the mask is

Mv,u=\{0,if​u∈Anc​\(v\)∪\{v\},−∞,otherwise,M\_\{v,u\}=\\begin\{cases\}0,&\\text\{if \}u\\in\\mathrm\{Anc\}\(v\)\\cup\\\{v\\\},\\\\ \-\\infty,&\\text\{otherwise\},\\end\{cases\}\(5\)where prefix tokens are visible to all tree nodes\. The masked attention for nodevvis computed as

Attn​\(Qv,K,V\)=softmax​\(Qv​K⊤d\+Mv\)​V\.\\mathrm\{Attn\}\(Q\_\{v\},K,V\)=\\mathrm\{softmax\}\\left\(\\frac\{Q\_\{v\}K^\{\\top\}\}\{\\sqrt\{d\}\}\+M\_\{v\}\\right\)V\.\(6\)This mask allows all tree nodes to be processed in parallel while ensuring that each branch follows an autoregressive\-like dependency structure\. More concretely, the tree\-causal mask induces a branch\-wise draft factorization

q​\(π​\(v\)∣x\)=∏u∈π​\(v\)q​\(yu∣x,hxo,π<u\),q\(\\pi\(v\)\\mid x\)=\\prod\_\{u\\in\\pi\(v\)\}q\(y\_\{u\}\\mid x,h\_\{x\}^\{o\},\\pi\_\{<u\}\),\(7\)This branch\-wise factorization mirrors the target model’s autoregressive factorization in Equation[4](https://arxiv.org/html/2606.18394#S2.E4), but conditions each draft node on the concrete ancestor tokens along its own path\. Compared with branch\-agnostic block\-diffusion drafting in Equation[3](https://arxiv.org/html/2606.18394#S2.E3), this makes the draft distribution more aligned with the autoregressiveMpM\_\{p\}while still allowing logits from all tree depths to be computed in parallel\.

### 2\.3Training

##### Data Preparation\.

We train the draft head on target\-aligned sequences, obtained either from the training corpus, as in multi\-token prediction\(Gloeckleet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib140); DeepSeek\-AI,[2024](https://arxiv.org/html/2606.18394#bib.bib149)\), or from continuations regenerated by the target model\(Agarwalet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib154); Guet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib170); Lianget al\.,[2026](https://arxiv.org/html/2606.18394#bib.bib171)\)\. Given a sequence\(x,y1:k\)\(x,y\_\{1:k\}\), we sample anchor positions and, at each anchor, construct a length\-NNtraining block ofNNconsecutive future positions\. The draft head predicts theseNNpositions in parallel under a block\-causal mask, matching the per\-position teacher logits produced by the target model on the same ground\-truth prefix\.

##### Distillation Loss\.

We train the draft head with a soft\-label distillation objective so that the draft distribution preserves the target model’s relative preferences across multiple plausible continuations\. Among distillation objectives, forward KL encourages the draft distribution to cover the target model’s probability mass, preserving soft\-label information beyond the top prediction\. Reverse KL, in contrast, is more mode\-seeking and can concentrate probability on high\-confidence teacher modes\(Guet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib170); Agarwalet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib154)\)\. We find this matters in practice \(Section[3](https://arxiv.org/html/2606.18394#S3)\): reverse KL underperforms substantially, while forward KL slightly outperforms hard\-label SFT\. We therefore adopt forward KL by default, and present the loss in this form below\.

For each active draft positionmm, letzq\(m\)z\_\{q\}^\{\(m\)\}andzp\(m\)z\_\{p\}^\{\(m\)\}denote the draft and target logits over the vocabulary𝒱\\mathcal\{V\}\. We define temperature\-normalized distributionsq~\(m\)=softmax​\(zq\(m\)/TKD\)\\tilde\{q\}^\{\(m\)\}=\\mathrm\{softmax\}\(z\_\{q\}^\{\(m\)\}/T\_\{\\mathrm\{KD\}\}\)andp~\(m\)=softmax​\(zp\(m\)/TKD\)\\tilde\{p\}^\{\(m\)\}=\\mathrm\{softmax\}\(z\_\{p\}^\{\(m\)\}/T\_\{\\mathrm\{KD\}\}\), and minimize the forward KL

ℒFKL\(m\)=DKL\(p~\(m\)∥q~\(m\)\)\.\\mathcal\{L\}\_\{\\mathrm\{FKL\}\}^\{\(m\)\}=D\_\{\\mathrm\{KL\}\}\\left\(\\tilde\{p\}^\{\(m\)\}\\,\\middle\\\|\\,\\tilde\{q\}^\{\(m\)\}\\right\)\.\(8\)The final objective is computed over all active draft\-token positions induced from the training sequence and normalized by the active\-position mask \(see Appendix[D](https://arxiv.org/html/2606.18394#A4)for details\):

ℒtrain=TKD2​∑mwm​ℒFKL\(m\)∑mwm,\\mathcal\{L\}\_\{\\mathrm\{train\}\}=T\_\{\\mathrm\{KD\}\}^\{2\}\\frac\{\\sum\_\{m\}w\_\{m\}\\mathcal\{L\}\_\{\\mathrm\{FKL\}\}^\{\(m\)\}\}\{\\sum\_\{m\}w\_\{m\}\},\(9\)wheremmindexes draft\-token positions within sampled blocks, excluding non\-loss and anchor\-token positions, andTKD\>0T\_\{\\mathrm\{KD\}\}\>0is the distillation temperature\. This objective preserves the target model’s soft\-label preferences and provides teacher\-aligned candidate distributions for tree expansion\.

### 2\.4Drafting and Verification

##### Parallel Tree Drafting\.

An overview of the design is shown in Figure[3](https://arxiv.org/html/2606.18394#S2.F3)\. At every decoding stepii, the causal draft head conditions on the hidden stateshxoh\_\{x\}^\{o\}of tokens that were verified and accepted in the previous step, and constructs a candidate tree with maximum draft depthNN, branching widthWW, and total node budgetBB\. With one draft\-head forward pass, we obtain logits for all draft depths and extract the top\-WWcandidate tokens at each depth\. Tree expansion is then driven by a branch scoring functions​\(π​\(v\)\)s\(\\pi\(v\)\), whose form can vary across implementations; we ablate different choices in Section[3\.4](https://arxiv.org/html/2606.18394#S3.SS4)\. By default, we use accumulated draft log\-probability for scoring:

s​\(π​\(v\)\)=∑u∈π​\(v\)log⁡q​\(yu∣x,hxo,π<u\),s\(\\pi\(v\)\)=\\sum\_\{u\\in\\pi\(v\)\}\\log q\(y\_\{u\}\\mid x,h\_\{x\}^\{o\},\\pi\_\{<u\}\),\(10\)whereπ<u\\pi\_\{<u\}denotes the ancestor tokens before nodeuu\. We then repeatedly pop the highest\-scoring expandable node, expand it with up toWWchildren, and reinsert the children with updated scores until the budget is filled or no expandable node remains\. This yields a candidate tree𝒯​\(x\)=\{π​\(v\)∣v∈𝒱𝒯\}\\mathcal\{T\}\(x\)=\\\{\\pi\(v\)\\mid v\\in\\mathcal\{V\}\_\{\\mathcal\{T\}\}\\\}, where each pathπ​\(v\)\\pi\(v\)is a speculative continuation; details are given in Algorithm[1](https://arxiv.org/html/2606.18394#alg1)\.

##### Tree Verification\.

After constructing the candidate tree𝒯​\(x\)\\mathcal\{T\}\(x\), the target modelMpM\_\{p\}verifies all tree nodes in parallel using tree attention\. For any candidate branchπ​\(v\)=y1:k\\pi\(v\)=y\_\{1:k\}, the target distribution factorizes autoregressively in compliance with Equation[4](https://arxiv.org/html/2606.18394#S2.E4)\.

The verifier then applies the speculative decoding acceptance rule along each candidate branch\. For each draft tokenyty\_\{t\}, we write the acceptance decision as

At∼Bernoulli\(αt\),αt=α\(yt;q\(⋅∣x,y<t\),p\(⋅∣x,y<t\)\),A\_\{t\}\\sim\\mathrm\{Bernoulli\}\(\\alpha\_\{t\}\),\\qquad\\alpha\_\{t\}=\\alpha\\\!\\left\(y\_\{t\};\\,q\(\\cdot\\mid x,y\_\{<t\}\),p\(\\cdot\\mid x,y\_\{<t\}\)\\right\),\(11\)whereqqis the draft distribution induced byMqM\_\{q\}, andα​\(⋅\)\\alpha\(\\cdot\)denotes the speculative verification rule\. In standard non\-greedy speculative sampling, this acceptance rule uses rejection sampling:

αt=min⁡\(1,p​\(yt∣x,y<t\)q​\(yt∣x,y<t\)\),\\alpha\_\{t\}=\\min\\\!\\left\(1,\\,\\frac\{p\(y\_\{t\}\\mid x,y\_\{<t\}\)\}\{q\(y\_\{t\}\\mid x,y\_\{<t\}\)\}\\right\),\(12\)with a target\-model correction token sampled when rejection occurs\(Leviathanet al\.,[2023](https://arxiv.org/html/2606.18394#bib.bib141)\)\. The accepted prefix length isa=max⁡\{r≤k:At=1,∀t≤r\}a=\\max\\left\\\{r\\leq k:A\_\{t\}=1,\\ \\forall t\\leq r\\right\\\}\. In the greedy setting,AtA\_\{t\}becomes deterministic: a draft token is accepted if it matches the target model’s next\-token prediction under the same context\.

## 3Experiments

### 3\.1Experiment Setup

##### Models and Datasets\.

We evaluateJetFlowon Qwen3\-8B and Qwen3\-30B\-A3B, covering both dense and MoE target models\(Yanget al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib172)\)\. Qwen3 supports both thinking and non\-thinking modes, and we use the non\-thinking mode throughout our main evaluation for efficient decoding\. For training data, we curate 780K examples from the Nemotron Post\-Training Dataset V2\(NVIDIA,[2025](https://arxiv.org/html/2606.18394#bib.bib173)\), including all available coding and math splits, random samples from STEM and chat splits, and 20K additional examples from CodeAlpaca\(Chaudhary,[2023](https://arxiv.org/html/2606.18394#bib.bib174)\)\. For regenerated training sequences, we apply the corresponding chat template to each data type and continue generation with the target model\. We evaluate on math benchmarks including GSM8K\(Cobbeet al\.,[2021](https://arxiv.org/html/2606.18394#bib.bib23)\), MATH\-500\(Hendryckset al\.,[2021](https://arxiv.org/html/2606.18394#bib.bib176)\), and AIME25; coding benchmarks including HumanEval\(Chenet al\.,[2021](https://arxiv.org/html/2606.18394#bib.bib177)\), MBPP\(Austinet al\.,[2021](https://arxiv.org/html/2606.18394#bib.bib178)\), and LiveCodeBench\(Jainet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib179)\); and open\-ended conversational tasks including MT\-Bench\(Zhenget al\.,[2023](https://arxiv.org/html/2606.18394#bib.bib33)\)\. All ablation studies are conducted on the math split of Nemotron Post\-Training Dataset V2\.

##### Baselines and Training Settings\.

We compare against EAGLE\-3 and DFlash, two strong head\-based speculative decoding baselines\(Liet al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib144); Chenet al\.,[2026](https://arxiv.org/html/2606.18394#bib.bib145)\)\. EAGLE\-3 uses multi\-layer feature fusion for direct token prediction, while DFlash uses a bidirectional block\-parallel draft head to generate multiple draft tokens in one pass\. DFlash\-T constructs draft trees from DFlash’s per\-position block\-diffusion distributions using best\-first tree expansion, resembles DDTree implementation, but uses a pure block\-diffusion head trained without loss weighting\(Ringel and Romano,[2026](https://arxiv.org/html/2606.18394#bib.bib192)\)\. For fair comparison, we trainJetFlowand all baselines on the same data mixture using a learning rate of3×10−43\\times 10^\{\-4\}on 8 H100 GPUs\. We also ablate DFlash\-style loss weighting in Section[3\.4](https://arxiv.org/html/2606.18394#S3.SS4); additional training details are provided in Appendix[D](https://arxiv.org/html/2606.18394#A4)\.

##### Implementation Details\.

For offline inference evaluation, we run experiments on 8 H100 or B200 GPUs\. Unless otherwise specified, results in Table[2](https://arxiv.org/html/2606.18394#S3.T2)use optimized Triton kernels for decoding\. For serving\-engine evaluation, we integrateJetFlowinto vLLM \(Appendix[E](https://arxiv.org/html/2606.18394#A5)\), part of which involve tree verification with a custom SM90 paged FlashAttention kernel using NVIDIA CuTe DSL, extending the paged attention forward path with shared\-memory tree\-mask staging and tree\-tail masking\.

### 3\.2Results

We evaluateJetFlowunder both greedy and non\-greedy decoding settings in compliance with Equation[12](https://arxiv.org/html/2606.18394#S2.E12)\. Evaluation results and comparisons with baselines are summarized in Table[1](https://arxiv.org/html/2606.18394#S3.T1)and Table[2](https://arxiv.org/html/2606.18394#S3.T2), and it showsJetFlowconsistently outperforms tree\-based baselines across all benchmarks\.

##### Low\-budget Regime\.

Compared with the autoregressive drafting baseline EAGLE\-3,JetFlowachieves substantially higher speedup at every reported budget by avoiding EAGLE\-3’s higher sequential drafting overhead\. At budget 16, DFlash andJetFlowperform nearly the same, suggesting that a short linear draft is already sufficient to cover many high\-probability continuations\. However, when increasing the budget to 32,JetFlowbrings slight improvement, while DFlash often saturates or degrades\. This indicates thatJetFlowconverts the additional draft budget into accepted tokens more effectively\.

Table 1:Low\-budget regime comparison ofJetFlowand baselines trained with the same Qwen3\-8B model and data recipe\. We report results on math, coding, and chat benchmarks using non\-thinking mode with a 3072 max tokens\. We report end\-to\-end decoding speedup over standard AR decoding and average accepted lengthτ\\tauon H100 GPU\.
##### High\-budget Regime\.

In Table[2](https://arxiv.org/html/2606.18394#S3.T2), as the draft budget increases,JetFlowshows a larger gain, e\.g\. improving from roughly44–6×6\\timesto77–10×10\\timesspeedup on math and coding benchmarks\. Whereas DFlash\-T scales more modestly to up to9×9\\times, showing causal tree drafting yields higher acceptance than branch\-agnostic tree construction\. And detailed qualitative example in Appendix[A\.2](https://arxiv.org/html/2606.18394#A1.SS2)on how causality helps\. Under temperature11,JetFlowremains effective, showing that the benefit of causal tree drafting remains robust under nongreedy decoding\.

Table 2:High\-budget comparison on Qwen3\-8B with at least 64 draft tokens\. EAGLE\-3 uses tree mode with max depth 8; larger budgets give minimal or worse gains due to training mismatch\.

### 3\.3System Performance

We evaluateJetFlowin an end\-to\-end serving setting by integrating tree drafting and verification into vLLM\. Serving performance depends not only on accepted length, but also on how the tree budget interacts with batching, verification overhead, memory traffic, and GPU occupancy\. Larger budgets can reduce the number of target\-model verification rounds, but also introduce higher draft\-head and tree\-verification cost, making the optimal budget load\-dependent\.

##### Budget Allocation\.

Table[11](https://arxiv.org/html/2606.18394#A5.T11)shows that larger tree budgets are most useful at small batch sizes\. At batch size 1, increasing the budget from 16 to 256 raises throughput from443\.3443\.3to968\.2968\.2TPS and speedup from3\.09×3\.09\\timesto6\.75×6\.75\\times\. As batch size increases, this gain diminishes: budget 256 drops to4\.51×4\.51\\timesat batch size 16 and2\.85×2\.85\\timesat batch size 32, nearly matching budget 128\. This indicates saturation, where the added verification cost and compute pressure reduce the benefit of accepting more tokens per step in throughput improvement\.

##### Practical implication\.

These results suggest that tree budgets should be selected according to serving load\. Large budgets are preferable in low\- to moderate\-load regimes, where spare GPU capacity can be used to reduce decoding iterations\. Under heavier load, smaller or moderate budgets may be more efficient by reducing per\-step overhead and preserving batching efficiency\. We evaluate static budget policies in this work and leave dynamic serving\-time budget scheduling to future work\.

### 3\.4Ablation Studies

Unless noted, all ablations in this section are evaluated on4×4\{\\times\}B200 with block size1616, FlashAttention\-2 for the target’s normal forward, Triton tree\-attention, and the accumulated draft log\-probability heuristic \(Equation[10](https://arxiv.org/html/2606.18394#S2.E10)\) for tree construction at node budget255255\. Reported speedup and average accepted lengthτ\\tauare invariant to hardware under lossless computation\.

#### 3\.4\.1Training Settings

##### Learning Rates\.

Table[3](https://arxiv.org/html/2606.18394#S3.T3)reports a grid search over learning rates from5×10−55\\times 10^\{\-5\}to1×10−31\\times 10^\{\-3\}\. Very small learning rates underperform and performance plateaus at3×10−43\\times 10^\{\-4\}\(peak speedup8\.30×8\.30\\times\), with6×10−46\\times 10^\{\-4\}and1×10−31\\times 10^\{\-3\}within∼2%\{\\sim\}2\\%of the peak\.

Table 3:Learning\-rate ablation withJetFlowand no loss weighting \(γ=0\\gamma=0; see Section[3\.4\.2](https://arxiv.org/html/2606.18394#S3.SS4.SSS2)forγ\\gamma’s definition and ablation\)\. We report speedup and average accepted lengthτ\\tau\.
##### Loss\.

Table[4](https://arxiv.org/html/2606.18394#S3.T4)compares SFT, forward\-KL distillation, and reverse\-KL distillation across four math benchmarks\. SFT and forward\-KL distillation match within∼3%\{\\sim\}3\\%at every dataset, while reverse\-KL distillation causes a3636–46%46\\%relative drop in comparison with forward\-KL’s, indicating the mode\-seeking objective concentrates draft probability mass too aggressively for tree drafting\.

Table 4:Loss\-objective ablation withJetFlowat LR6×10−46\{\\times\}10^\{\-4\}andγ=0\\gamma=0\. We report speedup and average accepted lengthτ\\tau\. All cells use checkpoints from a single training pipeline trained on math\-only data \(∼\\sim3 epochs\); we report on math benchmarks to keep the comparison in\-distribution\.
##### Model Generalizability\.

Table[5](https://arxiv.org/html/2606.18394#S3.T5)comparesJetFlowand DFlash\-T trained on Qwen3\-30B\-A3B, a MoE target model, with the same 800k data recipe\.JetFlowmaintains a competitively higher speedups across benchmarks, showing that causal tree drafting generalizes beyond dense Qwen3\-8B\.

Table 5:Model generalizability:JetFlowvs\. DFlash\-T on Qwen3\-30B\-A3B \(MoE target\), both trained with SFT on the same 800K\-example data mixture as our Qwen3\-8B main results\. Each cell reports speedup / average accepted lengthτ\\tauat temperature0with tree budget256256\.
##### Training Data\.

Table[6](https://arxiv.org/html/2606.18394#S3.T6)studies the effect of using different training data\. Regenerated target\-model sequences provide the strongest performance, confirming that matching the draft head to the target model’s own generations is important for maximizing accepted length\. In contrast,JetFlow\-Corpus trains directly on corpus sequences, it underperforms fully regeneratedJetFlow, but still achieves consistent speedups across all benchmarks, suggesting that causal parallel drafting can be incorporated during mid\-training or even pretraining, where regeneration would be prohibitively expensive\.

Table 6:Training\-data ablation:JetFlowvs\.JetFlow\-Corpus, both trained with SFT on the same 800K\-example data mixture \(Qwen3\-8B target\)\.JetFlowuses model\-regenerated continuations as supervision targets;JetFlow\-Corpus uses the original training corpus\. Each cell reports speedup / average accepted lengthτ\\tauat temperature0\.

#### 3\.4\.2Tree Drafting with Diffusion Head

Table[7](https://arxiv.org/html/2606.18394#S3.T7)compares causal and diffusion heads under different choices ofγ\\gamma, the parameter that controls how aggressively the DFlash training objective downweights per\-position loss at positions far from each anchor token\. Specifically, positioniiwithin a block contributes to the training loss with weightwi=exp⁡\(−max⁡\(i−ianchor,0\)/γ\)w\_\{i\}=\\exp\\\!\\left\(\-\\max\(i\-i\_\{\\mathrm\{anchor\}\},\\,0\)/\\gamma\\right\), whereγ=0\\gamma=0is interpreted as the limit of uniform per\-position weighting \(no decay\)\. The causal head is robust acrossγ\\gamma, while the diffusion head is sensitive toγ\\gammaand collapses at the endpoints \(from8\.36×8\.36\\timesatγ=7\\gamma=7to5\.46×5\.46\\timesatγ=0\\gamma=0and6\.17×6\.17\\timesatγ=15\\gamma=15\)\.

Table 7:Architecture andγ\\gammaablation on MATH\-500 withJetFlowat LR3×10−43\{\\times\}10^\{\-4\}with forward\-KL distillation\. We report speedup and average accepted lengthτ\\tau\.##### A failure example\.

On MATH\-500 prompt0at decode step0, the diffusion\-head draft \(γ=0\\gamma=0\) ranks “given told that” first in its tree, with cumulative draft surrogate∑ilog⁡ri=−3\.76\\sum\_\{i\}\\log r\_\{i\}=\-3\.76but cumulative target conditional∑ilog⁡p​\(yi∣x,π<vi\)=−63\.32\\sum\_\{i\}\\log p\(y\_\{i\}\\mid x,\\pi\_\{<v\_\{i\}\}\)=\-63\.32nats: the branch combines two mutually exclusive openers, exemplifying the failure mode ofqsurq\_\{\\mathrm\{sur\}\}in Equation[3](https://arxiv.org/html/2606.18394#S2.E3)\. The actually\-coherent “are given that the” \(target joint−0\.08\-0\.08\) sits at rank 3\. The causal\-head draft on the same prompt ranks “are told that” first, with surrogate≈\{\\approx\}target joint \(gap−0\.34\-0\.34\); tree verification accepts66tokens from the causal\-head tree but only44from the diffusion\-head tree \(visualized in Figure[4](https://arxiv.org/html/2606.18394#A1.F4)of Appendix[A](https://arxiv.org/html/2606.18394#A1)\)\.

##### Aggregate generalization\.

Across MATH\-500 prompts0–4949, the diffusion\-head’s rank\-11gap exceeds the causal\-head’s on92%92\\%of prompts, with a median5×5\\timeslarger \(\+62\.81\+62\.81vs\.\+12\.36\+12\.36nats\); the mean per\-step accepted\-token ratio is1\.95×1\.95\\times, closely matching the macroscopic1\.52×1\.52\\timesspeedup ratio in Table[7](https://arxiv.org/html/2606.18394#S3.T7)\. Full branch\-level tables and the rank\-11gap distribution are in Appendix[A](https://arxiv.org/html/2606.18394#A1)\.

## 4Related Work

##### Speculative Decoding\.

Speculative decoding accelerates autoregressive generation by drafting tokens and verifying them with the target model\(Leviathanet al\.,[2023](https://arxiv.org/html/2606.18394#bib.bib141); Chenet al\.,[2023](https://arxiv.org/html/2606.18394#bib.bib156)\)\. Its speedup depends on draft quality, which controls accepted length, and drafting cost, which controls wall\-clock gain\. Prior work improves these factors via draft\-target alignment\(Zhouet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib165); Liuet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib164); Goelet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib166)\), head\-based drafters such as Medusa and EAGLE\(Caiet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib142); Liet al\.,[2024b](https://arxiv.org/html/2606.18394#bib.bib143),[2025](https://arxiv.org/html/2606.18394#bib.bib144)\), or draft\-free retrieval methods such as Prompt Lookup Decoding, REST, PLD\+, SAM Decoding, and SuffixDecoding\(Saxena,[2023](https://arxiv.org/html/2606.18394#bib.bib180); Heet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib181); Somasundaramet al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib182); Huet al\.,[2025b](https://arxiv.org/html/2606.18394#bib.bib183); Oliaroet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib184)\)\. While these methods reduce overhead, they often rely on lexical overlap, retrieval quality, or repeated patterns\. In contrast,JetFlowtargets low\-cost parallel drafting with a learned causal head that preserves branch\-wise dependencies for high\-acceptance tree verification\.

##### Drafter Alignment\.

Previous works train draft models or draft heads to better match the target model\. Multi\-token prediction objectives supervise future\-token predictions during pretraining or midtraining\(Gloeckleet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib140); DeepSeek\-AI,[2024](https://arxiv.org/html/2606.18394#bib.bib149)\), while post\-training approaches use supervised fine\-tuning, offline distillation, online distillation, or direct acceptance\-rate optimization to improve agreement between draft and target distributions\(Zhouet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib165); Liuet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib164); Goelet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib166)\)\. These methods primarily improve the alignment term that controls acceptance rate\.JetFlowis complementary: it trains the drafter not only to match the target model, but also to preserve the causal structure needed for budgeted tree construction, so each branch is expanded from continuation\-aware draft distributions\.

##### Parallel Decoding and Drafting\.

Parallel decoding methods reduce latency by relaxing strict one\-token\-at\-a\-time generation\. Self\-speculative decoding drafts with early target\-model layers and verifies with later layers\(Elhoushiet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib185)\)\. Diffusion LLMs refine multiple tokens per iteration\(Wuet al\.,[2025](https://arxiv.org/html/2606.18394#bib.bib148); Qianet al\.,[2026](https://arxiv.org/html/2606.18394#bib.bib190)\)\. Jacobi\-style methods preserve causality by parallel fixed\-point updates, with recent consistency distillation works improving convergence to autoregressive outputs\(Fuet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib191); Kouet al\.,[2024](https://arxiv.org/html/2606.18394#bib.bib187); Huet al\.,[2025a](https://arxiv.org/html/2606.18394#bib.bib189)\)\. These methods increase tokens committed per forward, but are typically applied to the target model itself\. DFlash is the closest baseline, generating multiple draft tokens in one block\-parallel pass\(Chenet al\.,[2026](https://arxiv.org/html/2606.18394#bib.bib145)\); however, its branch\-agnostic block\-diffusion predictions are not explicitly optimized for causal tree construction\.

## 5Conclusion

We presentedJetFlow, a causal parallel tree drafting framework for speculative decoding that jointly addresses the two factors limiting end\-to\-end speedup: drafting latency and accepted length\. By preserving branch\-wise causal dependencies while generating candidate trees in a single low\-cost draft\-head pass,JetFlowmakes tree construction better aligned with the target model’s autoregressive factorization\. Across dense and MoE Qwen3 models and vLLM serving settings,JetFlowconsistently turns larger draft budgets into higher wall\-clock speedup, for MATH\-500, reaching9\.64×9\.64\\timeson Qwen3\-8B,9\.45×9\.45\\timeson Qwen3\-30B\-A3B, and over6×6\\timesserving speedup on a single H100 under low\-to\-moderate load\. These results show that block\-level causality conditioning over prefix hidden states is key to scaling speculative decoding with parallel drafting that is both cheap and effective\.

## References

- On\-policy distillation of language models: learning from self\-generated mistakes\.InInternational Conference on Learning Representations,Cited by:[§2\.3](https://arxiv.org/html/2606.18394#S2.SS3.SSS0.Px1.p1.4),[§2\.3](https://arxiv.org/html/2606.18394#S2.SS3.SSS0.Px2.p1.1)\.
- J\. Austin, A\. Odena, M\. Nye, M\. Bosma, H\. Michalewski, D\. Dohan, E\. Jiang, C\. Cai, M\. Terry, Q\. V\. Le, and C\. Sutton \(2021\)Program synthesis with large language models\.arXiv preprint arXiv:2108\.07732\.External Links:2108\.07732Cited by:[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px1.p1.1)\.
- T\. Cai, Y\. Li, Z\. Geng, H\. Peng, J\. D\. Lee, D\. Chen, and T\. Dao \(2024\)Medusa: simple llm inference acceleration framework with multiple decoding heads\.arXiv preprint arXiv:2401\.10774\.External Links:[Link](https://arxiv.org/abs/2401.10774)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p3.1),[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px2.p1.3),[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px2.p2.3),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1)\.
- S\. Chaudhary \(2023\)Code Alpaca: an instruction\-following LLaMA model trained on code generation instructions\.GitHub\.Note:[https://github\.com/sahil280114/codealpaca](https://github.com/sahil280114/codealpaca)Cited by:[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px1.p1.1)\.
- C\. Chen, S\. Borgeaud, G\. Irving, J\. Lespiau, L\. Sifre, and J\. Jumper \(2023\)Accelerating large language model decoding with speculative sampling\.arXiv preprint arXiv:2302\.01318\.External Links:[Link](https://arxiv.org/abs/2302.01318)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p2.1),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1)\.
- J\. Chen, Y\. Liang, and Z\. Liu \(2026\)DFlash: block diffusion for flash speculative decoding\.arXiv preprint arXiv:2602\.06036\.External Links:[Link](https://arxiv.org/abs/2602.06036)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p3.1),[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px2.p2.3),[§2\.2](https://arxiv.org/html/2606.18394#S2.SS2.SSS0.Px3.p1.1),[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px2.p1.1),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px3.p1.1)\.
- M\. Chen, J\. Tworek, H\. Jun, Q\. Yuan, H\. P\. d\. O\. Pinto, J\. Kaplan, H\. Edwards, Y\. Burda, N\. Joseph, G\. Brockman, A\. Ray, R\. Puri, G\. Krueger, M\. Petrov, H\. Khlaaf, G\. Sastry, P\. Mishkin, B\. Chan, S\. Gray, N\. Ryder, M\. Pavlov, A\. Power, L\. Kaiser, M\. Bavarian, C\. Winter, P\. Tillet, F\. P\. Such, D\. Cummings, M\. Plappert, F\. Chantzis, E\. Barnes, A\. Herbert\-Voss, W\. H\. Guss, A\. Nichol, A\. Paino, N\. Tezak, J\. Tang, I\. Babuschkin, S\. Balaji, S\. Jain, W\. Saunders, C\. Hesse, A\. N\. Carr, J\. Leike, J\. Achiam, V\. Misra, E\. Morikawa, A\. Radford, M\. Knight, M\. Brundage, M\. Murati, K\. Mayer, P\. Welinder, B\. McGrew, D\. Amodei, S\. McCandlish, I\. Sutskever, and W\. Zaremba \(2021\)Evaluating large language models trained on code\.arXiv preprint arXiv:2107\.03374\.External Links:2107\.03374Cited by:[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px1.p1.1)\.
- K\. Cobbe, V\. Kosaraju, M\. Bavarian, M\. Chen, H\. Jun, L\. Kaiser, M\. Plappert, J\. Tworek, J\. Hilton, R\. Nakano, C\. Hesse, and J\. Schulman \(2021\)Training verifiers to solve math word problems\.arXiv preprint arXiv:2110\.14168\.Cited by:[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px1.p1.1)\.
- DeepSeek\-AI \(2024\)DeepSeek\-v3 technical report\.arXiv preprint arXiv:2412\.19437\.External Links:[Link](https://arxiv.org/abs/2412.19437)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p2.1),[§2\.3](https://arxiv.org/html/2606.18394#S2.SS3.SSS0.Px1.p1.4),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px2.p1.1)\.
- DeepSeek\-AI \(2026\)DeepSeek\-v4: towards highly efficient million\-token context intelligence with multi\-head latent attention, manifold\-constrained hyper\-connections and multi\-token prediction\.Technical report\.External Links:[Link](https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/resolve/main/DeepSeek_V4.pdf)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p2.1)\.
- X\. Denget al\.\(2025\)SWE\-bench pro: can ai agents solve long\-horizon software engineering tasks?\.arXiv preprint arXiv:2509\.16941\.External Links:[Link](https://arxiv.org/abs/2509.16941)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p1.1)\.
- M\. Elhoushi, A\. Shrivastava, D\. Liskovich, B\. Hosmer, B\. Wasti, L\. Lai, A\. Mahmoud, B\. Acun, S\. Agarwal, A\. Roman, A\. Aly, B\. Chen, C\. Wu, C\. Li, C\. Wang, G\. Goswami, P\. Panda, R\. Tian, X\. Wu, Y\. Jia, K\. Keutzer, A\. Gholami, J\. S\. Lumbroso, and Z\. Chen \(2024\)LayerSkip: enabling early exit inference and self\-speculative decoding\.arXiv preprint arXiv:2404\.16710\.External Links:2404\.16710Cited by:[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px3.p1.1)\.
- Y\. Fu, P\. Bailis, I\. Stoica, and H\. Zhang \(2024\)Break the sequential dependency of LLM inference using lookahead decoding\.InProceedings of the 41st International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.235,pp\. 14060–14079\.Cited by:[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px3.p1.1)\.
- F\. Gloeckle, B\. Y\. Idrissi, B\. Rozière, D\. Lopez\-Paz, and G\. Synnaeve \(2024\)Better & faster large language models via multi\-token prediction\.arXiv preprint arXiv:2404\.19737\.External Links:[Link](https://arxiv.org/abs/2404.19737)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p2.1),[§2\.3](https://arxiv.org/html/2606.18394#S2.SS3.SSS0.Px1.p1.4),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px2.p1.1)\.
- R\. Goel, M\. Gagrani, W\. Jeon, J\. Park, M\. Lee, and C\. Lott \(2024\)Direct alignment of draft model for speculative decoding with chat\-fine\-tuned llms\.arXiv preprint arXiv:2403\.00858\.Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p3.1),[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px2.p1.3),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px2.p1.1)\.
- Y\. Gu, L\. Dong, F\. Wei, and M\. Huang \(2024\)MiniLLM: knowledge distillation of large language models\.InInternational Conference on Learning Representations,Cited by:[§2\.3](https://arxiv.org/html/2606.18394#S2.SS3.SSS0.Px1.p1.4),[§2\.3](https://arxiv.org/html/2606.18394#S2.SS3.SSS0.Px2.p1.1)\.
- Z\. He, Z\. Zhong, T\. Cai, J\. D\. Lee, and D\. He \(2024\)REST: retrieval\-based speculative decoding\.InProceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,pp\. 1582–1595\.Cited by:[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1)\.
- D\. Hendrycks, C\. Burns, S\. Kadavath, A\. Arora, S\. Basart, E\. Tang, D\. Song, and J\. Steinhardt \(2021\)Measuring mathematical problem solving with the MATH dataset\.InNeurIPS Datasets and Benchmarks,Cited by:[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px1.p1.1)\.
- L\. Hu, S\. Kou, Y\. Fu, S\. Rajbhandari, T\. Rosing, Y\. He, Z\. Deng, and H\. Zhang \(2025a\)Fast and accurate causal parallel decoding using jacobi forcing\.arXiv preprint arXiv:2512\.14681\.External Links:2512\.14681Cited by:[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px3.p1.1)\.
- Y\. Hu, Z\. Chen, D\. Zhang, Y\. Wang, Y\. Zhang, B\. Chen, and A\. Li \(2025b\)SAM decoding: speculative decoding via suffix automaton\.InProceedings of the 63rd Annual Meeting of the Association for Computational Linguistics,Cited by:[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1)\.
- A\. Huang, C\. Yao, C\. Wan, M\. Chen,et al\.\(2026\)Step 3\.5 flash: open frontier\-level intelligence with 11b active parameters\.arXiv preprint arXiv:2602\.10604\.External Links:[Link](https://arxiv.org/abs/2602.10604)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p2.1)\.
- N\. Jain, K\. Han, A\. Gu, W\. Li, F\. Yan, T\. Zhang, S\. Wang, A\. Solar\-Lezama, K\. Sen, and I\. Stoica \(2024\)LiveCodeBench: holistic and contamination free evaluation of large language models for code\.arXiv preprint arXiv:2403\.07974\.External Links:2403\.07974Cited by:[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px1.p1.1)\.
- N\. Jain, K\. Han, A\. Gu, W\. Li, F\. Yan, T\. Zhang, S\. Wang, A\. Solar\-Lezama, K\. Sen, and I\. Stoica \(2025\)LiveCodeBench: holistic and contamination free evaluation of large language models for code\.InThe Thirteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=chfJJYC3iL)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p1.1)\.
- S\. Kou, L\. Hu, Z\. He, Z\. Deng, and H\. Zhang \(2024\)CLLMs: consistency large language models\.InProceedings of the 41st International Conference on Machine Learning,Cited by:[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px3.p1.1)\.
- Y\. Leviathan, M\. Kalman, and Y\. Matias \(2023\)Fast inference from transformers via speculative decoding\.InProceedings of the 40th International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.202,pp\. 19274–19286\.External Links:[Link](https://arxiv.org/abs/2211.17192)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p2.1),[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px1.p1.7),[§2\.4](https://arxiv.org/html/2606.18394#S2.SS4.SSS0.Px2.p2.6),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1)\.
- Y\. Li, F\. Wei, C\. Zhang, and H\. Zhang \(2024a\)EAGLE\-2: faster inference of language models with dynamic draft trees\.arXiv preprint arXiv:2406\.16858\.External Links:[Link](https://arxiv.org/abs/2406.16858)Cited by:[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px2.p1.3)\.
- Y\. Li, F\. Wei, C\. Zhang, and H\. Zhang \(2024b\)EAGLE: speculative sampling requires rethinking feature uncertainty\.arXiv preprint arXiv:2401\.15077\.External Links:[Link](https://arxiv.org/abs/2401.15077)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p3.1),[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px2.p2.3),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1)\.
- Y\. Li, F\. Wei, C\. Zhang, and H\. Zhang \(2025\)EAGLE\-3: scaling up inference acceleration of large language models via training\-time test\.arXiv preprint arXiv:2503\.01840\.External Links:[Link](https://arxiv.org/abs/2503.01840)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p3.1),[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px2.p1.3),[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px2.p2.3),[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px2.p1.1),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1)\.
- Y\. Liang, Z\. Wang, H\. Chen, X\. Sun, J\. Wu, X\. Yu, J\. Liu, E\. Barsoum, Z\. Liu, and N\. K\. Jha \(2026\)CD4LM: consistency distillation and adaptive decoding for diffusion language models\.arXiv preprint arXiv:2601\.02236\.Cited by:[§2\.3](https://arxiv.org/html/2606.18394#S2.SS3.SSS0.Px1.p1.4)\.
- X\. Liu, L\. Hu, P\. Bailis, A\. Cheung, Z\. Deng, I\. Stoica, and H\. Zhang \(2024\)Online speculative decoding\.InProceedings of the 41st International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.235,pp\. 31131–31146\.Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p3.1),[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px2.p1.3),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px2.p1.1)\.
- J\. Lu, T\. Holleis, Y\. Zhang, B\. Aumayer, F\. Nan, F\. Bai, S\. Ma, S\. Ma, M\. Li, G\. Yin, Z\. Wang, and R\. Pang \(2025\)ToolSandbox: a stateful, conversational, interactive evaluation benchmark for llm tool use capabilities\.InFindings of the Association for Computational Linguistics: NAACL 2025,pp\. 1160–1183\.External Links:[Link](https://aclanthology.org/2025.findings-naacl.68/)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p1.1)\.
- X\. Miao, G\. Oliaro, Z\. Zhang, X\. Cheng, Z\. Wang, Z\. Zhang, R\. Y\. Y\. Wong, A\. Zhu, L\. Yang, X\. Shi, C\. Shi, Z\. Chen, D\. Arfeen, R\. Abhyankar, and Z\. Jia \(2024\)SpecInfer: accelerating generative large language model serving with tree\-based speculative inference and verification\.InProceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems,Cited by:[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px2.p1.3)\.
- NVIDIA \(2025\)Nemotron post\-training dataset v2\.Note:[https://huggingface\.co/datasets/nvidia/Nemotron\-Post\-Training\-Dataset\-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2)DatasetCited by:[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px1.p1.1)\.
- G\. Oliaro, Z\. Jia, D\. Campos, and A\. Qiao \(2024\)SuffixDecoding: a model\-free approach to speeding up large language model inference\.arXiv preprint arXiv:2411\.04975\.External Links:2411\.04975Cited by:[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1)\.
- A\. P\. Patil, V\. Jain, X\. Zhang, K\. Ganesan, S\. K\. Bird, J\. Zou, M\. Zaharia, I\. Stoica, and J\. E\. Gonzalez \(2025\)BFCL v3: evaluating large language models on real\-world function calling and multi\-turn tool use\.InProceedings of the 42nd International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.267,pp\. 38904–38941\.External Links:[Link](https://proceedings.mlr.press/v267/patil25a.html)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p1.1)\.
- Y\. Qian, J\. Su, L\. Hu, P\. Zhang, Z\. Deng, P\. Zhao, and H\. Zhang \(2026\)d3LLM: ultra\-fast diffusion LLM using pseudo\-trajectory distillation\.arXiv preprint arXiv:2601\.07568\.External Links:2601\.07568Cited by:[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px3.p1.1)\.
- L\. Ringel and Y\. Romano \(2026\)Accelerating speculative decoding with block diffusion draft trees\.arXiv preprint arXiv:2604\.12989\.Cited by:[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px2.p1.1)\.
- A\. Saxena \(2023\)Prompt lookup decoding\.Note:[https://github\.com/apoorvumang/prompt\-lookup\-decoding](https://github.com/apoorvumang/prompt-lookup-decoding)GitHub repositoryCited by:[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1)\.
- S\. Somasundaram, H\. Honnappa, X\. Zhang, and A\. Kalyan \(2025\)PLD\+: accelerating LLM inference by leveraging language model artifacts\.InFindings of the Association for Computational Linguistics: NAACL 2025,Cited by:[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1)\.
- J\. Wei, X\. Wang, D\. Schuurmans, M\. Bosma, B\. Ichter, F\. Xia, E\. H\. Chi, Q\. V\. Le, and D\. Zhou \(2022\)Chain\-of\-thought prompting elicits reasoning in large language models\.InAdvances in Neural Information Processing Systems,Vol\.35,pp\. 24824–24837\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2022/file/9d5609613524ecf4f15af0f7b31abca4-Paper-Conference.pdf)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p1.1)\.
- C\. Wu, H\. Zhang, S\. Xue, S\. Diao, Y\. Fu, Z\. Liu, P\. Molchanov, P\. Luo, S\. Han, and E\. Xie \(2025\)Fast\-dllm v2: efficient block\-diffusion llm\.arXiv preprint arXiv:2509\.26328\.External Links:[Link](https://arxiv.org/abs/2509.26328)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p2.1),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px3.p1.1)\.
- B\. Xiao, B\. Xia, B\. Yang, B\. Gao, B\. Shen, C\. Zhang, C\. He,et al\.\(2026\)MiMo\-v2\-flash technical report\.arXiv preprint arXiv:2601\.02780\.External Links:[Link](https://arxiv.org/abs/2601.02780)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p2.1)\.
- A\. Yang, A\. Li, B\. Yang, B\. Zhang, B\. Hui, B\. Zheng, B\. Yu, C\. Gao, C\. Huang, C\. Lv, C\. Zheng, D\. Liu, F\. Zhou, F\. Huang, F\. Hu, H\. Ge, H\. Wei, H\. Lin, J\. Tang, J\. Yang, J\. Tu, J\. Zhang, J\. Yang, J\. Yang, J\. Zhou, J\. Zhou, J\. Lin, K\. Dang, K\. Bao, K\. Yang, L\. Yu, L\. Deng, M\. Li, M\. Xue, M\. Li, P\. Zhang, P\. Wang, Q\. Zhu, R\. Men, R\. Gao, S\. Liu, S\. Luo, T\. Li, T\. Tang, W\. Yin, X\. Ren, X\. Wang, X\. Zhang, X\. Ren, Y\. Fan, Y\. Su, Y\. Zhang, Y\. Zhang, Y\. Wan, Y\. Liu, Z\. Wang, Z\. Cui, Z\. Zhang, Z\. Zhou, and Z\. Qiu \(2025\)Qwen3 technical report\.arXiv preprint arXiv:2505\.09388\.External Links:2505\.09388Cited by:[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px1.p1.1)\.
- S\. Yao, N\. Shinn, P\. Razavi, and K\. R\. Narasimhan \(2025\)τ\\tau\-bench: a benchmark for tool\-agent\-user interaction in real\-world domains\.InThe Thirteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=roNSXZpUDN)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p1.1)\.
- S\. Zhao, D\. Gupta, Q\. Zheng, and A\. Grover \(2025\)D1: scaling reasoning in diffusion large language models via reinforcement learning\.arXiv preprint arXiv:2504\.12216\.External Links:[Link](https://arxiv.org/abs/2504.12216)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p2.1)\.
- L\. Zheng, W\. Chiang, Y\. Sheng, S\. Zhuang, Z\. Wu, Y\. Zhuang, Z\. Lin, Z\. Li, D\. Li, E\. Xing,et al\.\(2023\)Judging llm\-as\-a\-judge with mt\-bench and chatbot arena\.Advances in Neural Information Processing Systems36,pp\. 46595–46623\.Cited by:[§3\.1](https://arxiv.org/html/2606.18394#S3.SS1.SSS0.Px1.p1.1)\.
- Y\. Zhou, K\. Lyu, A\. S\. Rawat, A\. K\. Menon, A\. Rostamizadeh, S\. Kumar, J\. Kagy, and R\. Agarwal \(2024\)DistillSpec: improving speculative decoding via knowledge distillation\.InInternational Conference on Learning Representations,Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p3.1),[§2\.1](https://arxiv.org/html/2606.18394#S2.SS1.SSS0.Px2.p1.3),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px1.p1.1),[§4](https://arxiv.org/html/2606.18394#S4.SS0.SSS0.Px2.p1.1)\.
- D\. Zhu, X\. Wei, G\. Zhao, W\. Wu, H\. Zou, J\. Ran, XWang, L\. Sun, X\. Zhang, and S\. Li \(2025\)Chain\-of\-thought matters: improving long\-context language models with reasoning path supervision\.InFindings of the Association for Computational Linguistics: EMNLP 2025,pp\. 3197–3211\.External Links:[Link](https://aclanthology.org/2025.findings-emnlp.170/)Cited by:[§1](https://arxiv.org/html/2606.18394#S1.p1.1)\.

## Appendix ATree Drafting Quality Analysis

This appendix provides the full data behind the rank\-11case study and theN=50N=50aggregate statistics summarized in Section[3\.4\.2](https://arxiv.org/html/2606.18394#S3.SS4.SSS2)\. Figure[4](https://arxiv.org/html/2606.18394#A1.F4)visualizes the rank\-11\+ rank\-33contrast for both heads on MATH\-500 prompt\#​0\\\#0, decode step0\(root token‘‘We’’\); the full top\-55branches and the rank\-11gap distribution follow\.

\{tikzpicture\}

\[ tok/\.style= rectangle, draw, rounded corners=2pt, inner sep=4pt, font=, text height=1\.6ex, text depth=0\.4ex, faithful/\.style=tok, fill=green\!10, draw=green\!55\!black, incoherent/\.style=tok, fill=red\!8, draw=red\!65\!black, rowlabel/\.style=font=, anchor=east, inner sep=2pt, panelhdr/\.style=font=, anchor=west, annot/\.style=font=, anchor=west, inner sep=2pt\]

\\node

\[panelhdr\] at \(0, 4\.6\) \(a\) Causal head,γ=0\\gamma=0;\\node\[rowlabel\] at \(1\.7, 3\.8\) rank 1✓\\checkmark;\\node\[faithful, anchor=west\] at \(1\.75, 3\.8\) \(c1a\) are;\\node\[faithful, anchor=west, xshift=1\.5mm\] at \(c1a\.east\) \(c1b\) told;\\node\[faithful, anchor=west, xshift=1\.5mm\] at \(c1b\.east\) \(c1c\) that;\\node\[annot, xshift=4mm\] at \(c1c\.east\) gap=−0\.34=\\mathbf\{\-0\.34\}faithful⇒\\Rightarrowverifier accepts6tokens;\\node\[rowlabel\] at \(1\.7, 2\.9\) rank 3×\\times;\\node\[incoherent, anchor=west\] at \(1\.75, 2\.9\) \(c3a\) are;\\node\[incoherent, anchor=west, xshift=1\.5mm\] at \(c3a\.east\) \(c3b\) given;\\node\[incoherent, anchor=west, xshift=1\.5mm\] at \(c3b\.east\) \(c3c\) that;\\node\[incoherent, anchor=west, xshift=1\.5mm\] at \(c3c\.east\) \(c3d\) the2product;\\node\[annot, xshift=4mm\] at \(c3d\.east\) gap=\+42\.50=\+42\.50;

\\node

\[panelhdr\] at \(0, 1\.7\) \(b\) Diffusion head,γ=0\\gamma=0;\\node\[rowlabel\] at \(1\.7, 0\.9\) rank 1×\\times;\\node\[incoherent, anchor=west\] at \(1\.75, 0\.9\) \(d1a\) given;\\node\[incoherent, anchor=west, xshift=1\.5mm\] at \(d1a\.east\) \(d1b\) told;\\node\[incoherent, anchor=west, xshift=1\.5mm\] at \(d1b\.east\) \(d1c\) that;\\node\[annot, xshift=4mm\] at \(d1c\.east\) gap=\+59\.56=\\mathbf\{\+59\.56\}incoherent⇒\\Rightarrowverifier accepts4tokens;\\node\[rowlabel\] at \(1\.7, 0\.0\) rank 3✓\\checkmark;\\node\[faithful, anchor=west\] at \(1\.75, 0\.0\) \(d3a\) are;\\node\[faithful, anchor=west, xshift=1\.5mm\] at \(d3a\.east\) \(d3b\) given;\\node\[faithful, anchor=west, xshift=1\.5mm\] at \(d3b\.east\) \(d3c\) that;\\node\[faithful, anchor=west, xshift=1\.5mm\] at \(d3c\.east\) \(d3d\) the;\\node\[annot, xshift=4mm\] at \(d3d\.east\) gap=−3\.69=\-3\.69faithful, but at rank 3;

Figure 4:Tree\-quality failure mode at MATH\-500 prompt\#​0\\\#0, decode step0\.Both heads draft from the same prefix \(last token‘‘We’’\)\. The causal head’s rank\-11branch \(“are told that”\) is faithful: target jointΣ​log⁡p≈Σ​log⁡r\\Sigma\\log p\\approx\\Sigma\\log r, so tree verification walks66tokens along it\. The diffusion head’s rank\-11branch \(“given told that”\) is incoherent \(target jointΣ​log⁡p=−63\.32\\Sigma\\log p=\-63\.32nats, i\.e\. probability≈e−63\{\\approx\}e^\{\-63\}\) because its branch\-agnostic per\-position predictor composes “given” \(depth11\) and “told” \(depth22\) independently in the surrogate, even though no real continuation places these two words consecutively \(the surrogateqsurq\_\{\\mathrm\{sur\}\}is formally defined in Equation[3](https://arxiv.org/html/2606.18394#S2.E3)\)\. The actually\-coherent “are given that the” is in the diffusion tree but only at rank33; the surrogate fails to promote it, and the verifier accepts only44tokens\.### A\.1Setup

We modify the speculative decoding loop to dump the constructed tree at a chosen decode step\. For every nodevvwe record the cumulative draft surrogate∑i=1depth​\(v\)log⁡ri​\(yvi∣x\)\\sum\_\{i=1\}^\{\\mathrm\{depth\}\(v\)\}\\log r\_\{i\}\(y\_\{v\_\{i\}\}\\mid x\)and the cumulative target conditional∑i=1depth​\(v\)log⁡p​\(yvi∣x,π<vi\)\\sum\_\{i=1\}^\{\\mathrm\{depth\}\(v\)\}\\log p\(y\_\{v\_\{i\}\}\\mid x,\\pi\_\{<v\_\{i\}\}\), along with the tree\-verifier’s accepted path\. The two drafts compared are a causal head and a diffusion head, both trained at LR3×10−43\{\\times\}10^\{\-4\}withγ=0\\gamma=0, evaluated on4×4\{\\times\}B200 with block sizeN=16N=16\(also the maximum draft depth\), FlashAttention\-2, and Triton tree\-attention\. Both heads share the prefill and produceNNparallel position\-wise distributionsr1,…,rNr\_\{1\},\\ldots,r\_\{N\}; tree construction uses the cumulative\-log\-probability heuristic with branching widthW=7W=7and node budgetB=255B=255\.

### A\.2Top\-5 Branches at the Canonical Example

Table[8](https://arxiv.org/html/2606.18394#A1.T8)extends Figure[4](https://arxiv.org/html/2606.18394#A1.F4)to the full top\-55branches of each head’s tree at MATH\-500 prompt\#​0\\\#0, decode step0\(root token‘‘We’’\)\. The pattern reported in the main text repeats throughout the tree\. For the diffusion head, top\-22and top\-44both combine “given” at depth11and “told” at depth22with target joints below−50\-50nats; only rank\-33\(target joint−0\.08\-0\.08\) is coherent\. For the causal head, the rank\-11branch \(“are told that”\) is faithful \(gap−0\.34\-0\.34\); ranks22–55all share the off\-argmax depth\-22token “given” \(vs\. rank\-11’s argmax “told”\) and degrade at depths44–55\(e\.g\. “the2product” at rank33, gap\+42\.50\+42\.50\), consistent with the off\-argmax inheritance described in §[A\.4](https://arxiv.org/html/2606.18394#A1.SS4): each branch’s depth\-ddmarginal was anchored to the argmax extension at depthd−1d\-1, so off\-argmax branches inherit a conditioning context that does not match their own ancestor token\.

Table 8:Top\-5 branches of each head’s tree at MATH\-500 prompt0, decode step0\. Both drafts run on the identical prefix with the same target\.Σ​log⁡r\\Sigma\\log randΣ​log⁡p\\Sigma\\log pin nats;Δ=\\Delta=surrogate−\-target\.
### A\.3Gap Distribution over5050Prompts

Table[9](https://arxiv.org/html/2606.18394#A1.T9)reports the rank\-11gap distribution across MATH\-500 prompts0–4949for both heads atγ=0\\gamma=0and atγ=7\\gamma=7\(DFlash’s best macroscopic loss\-weighting setting, Table[7](https://arxiv.org/html/2606.18394#S3.T7)\)\. Atγ=0\\gamma=0, the diffusion head exhibits a dramatic upper\-tail failure:26%26\\%of prompts have rank\-11gap≥\+80\\geq\+80nats \(joint target probability<10−35<10^\{\-35\}\), versus0%0\\%for the causal head; faithful rank\-11\(gap<\+5<\+5nats\) appears in only6%6\\%of diffusion prompts versus42%42\\%for causal\. The macroscopic effect is direct: mean accepted length is4\.844\.84tokens for diffusion vs9\.469\.46for causal atγ=0\\gamma=0\. Atγ=7\\gamma=7, the diffusion head’s tree quality recovers substantially — its extreme tail drops to4%4\\%and faithful rank\-11rises to26%26\\%— with mean acceptance reaching9\.429\.42tokens\. It does not, however, match the causal head, which retains the lower extreme tail \(2%2\\%\) and higher mean acceptance \(9\.649\.64\) atγ=7\\gamma=7\. The takeaway is that the causal head’s rank\-11faithfulness is robust acrossγ\\gamma, whereas the diffusion head depends on loss\-weighting tuning to recover from theγ=0\\gamma=0failure mode\.JetFlow’s contribution is the structural robustness this provides: the causal mask between depths obviates the need to tuneγ\\gamma\.

Table 9:Distribution of rank\-11surrogate–target gaps across MATH\-500 prompts0–4949\(each at the prompt’s natural step\-0tree\)\.γ=7\\gamma=7is DFlash’s best macroscopic loss\-weighting setting \(Table[7](https://arxiv.org/html/2606.18394#S3.T7)\)\. Mean accepted length per step is reported in the bottom row for context\.
### A\.4Why the Failure Mode Is Structural

The diffusion head produces allNNposition\-wise distributionsr1,…,rNr\_\{1\},\\ldots,r\_\{N\}from a single shared hidden state with no causal mask between depths, sor2r\_\{2\}does not condition ony1y\_\{1\}\. Whenr1r\_\{1\}andr2r\_\{2\}both place high mass on tokens that fit a given depth marginally but cannot follow each other syntactically \(e\.g\. “given” at depth11and “told” at depth22\), the surrogater1​\(y1\)⋅r2​\(y2\)r\_\{1\}\(y\_\{1\}\)\\cdot r\_\{2\}\(y\_\{2\}\)promotes their composition into the rank\-11branch\. The causal head’s parallel marginals are produced with causal masking between depths, sor2r\_\{2\}is anchored to the model’s own argmax at depth11; off\-argmax branches in the heap still inherit this anchoredr2r\_\{2\}, but the rank\-11branch \(the argmax\-aligned one, which dominates verifier behavior\) receives a faithful surrogate score\. This anchoring is what closes the gap betweenqsurq\_\{\\mathrm\{sur\}\}and the target’s branch\-conditional joint at the head of the tree, and it is what allows the causal head to remain robust acrossγ\\gamma: the rank\-11alignment does not depend on a loss\-weighting prior over depths\.

## Appendix BTree\-Drafting Algorithm Details

Algorithm[1](https://arxiv.org/html/2606.18394#alg1)describes howJetFlowconstructs a candidate tree under a fixed node budget\. Starting from the root prefixxx, the algorithm maintains a priority queue of expandable nodes ranked by a branch scoring functionScore​\(⋅\)\\mathrm\{Score\}\(\\cdot\)\. At each iteration, it expands the highest\-scoring node by adding up toWWchildren at the next depth, unless the node has already reached the maximum draft depthNN\. Each new child is inserted back into the queue with its updated path score, and the process continues until the total node budgetBBis exhausted\. The returned tree𝒯​\(x\)\\mathcal\{T\}\(x\)therefore contains the highest\-scoring candidate branches according to the chosen scoring rule\.

Algorithm 1Parallel Tree Drafting1:Prefix

xx, maximum draft depth

NN, branching width

WW, node budget

BB, scoring function

Score​\(⋅\)\\mathrm\{Score\}\(\\cdot\)
2:Candidate tree

𝒯​\(x\)\\mathcal\{T\}\(x\)
3:Initialize tree

𝒯\\mathcal\{T\}with root node

v0v\_\{0\}
4:Initialize priority queue

𝒬←\{\(v0,Score​\(π​\(v0\)\)\)\}\\mathcal\{Q\}\\leftarrow\\\{\(v\_\{0\},\\mathrm\{Score\}\(\\pi\(v\_\{0\}\)\)\)\\\}
5:while

\|𝒱𝒯\|<B\|\\mathcal\{V\}\_\{\\mathcal\{T\}\}\|<Band

𝒬≠∅\\mathcal\{Q\}\\neq\\emptysetdo

6:Pop the highest\-scoring node

vvfrom

𝒬\\mathcal\{Q\}
7:if

depth​\(v\)=N\\mathrm\{depth\}\(v\)=Nthen

8:continue

9:endif

10:Obtain up to

WWcandidate children

𝒞​\(v\)\\mathcal\{C\}\(v\)for the next depth

11:for

y∈𝒞​\(v\)y\\in\\mathcal\{C\}\(v\)do

12:if

\|𝒱𝒯\|=B\|\\mathcal\{V\}\_\{\\mathcal\{T\}\}\|=Bthen

13:break

14:endif

15:Add child node

uuwith token

yu=yy\_\{u\}=yand parent

vvto

𝒯\\mathcal\{T\}
16:Compute

su←Score​\(π​\(u\)\)s\_\{u\}\\leftarrow\\mathrm\{Score\}\(\\pi\(u\)\)
17:Push

\(u,su\)\(u,s\_\{u\}\)into

𝒬\\mathcal\{Q\}
18:endfor

19:endwhile

20:return

𝒯​\(x\)=\{π​\(v\)∣v∈𝒱𝒯\}\\mathcal\{T\}\(x\)=\\\{\\pi\(v\)\\mid v\\in\\mathcal\{V\}\_\{\\mathcal\{T\}\}\\\}

## Appendix CTree\-Drafting Algorithm Ablation

We compare three tree\-construction scoring algorithms at the production setting: cumulative draft log\-probability \(accum\_logp, our default\), entropy\-guided \(per\-depth entropy alone\), and hybrid \(cumulative log\-probability plusα\\alpha\-weighted per\-depth entropy\)\. Results are summarized in Table[10](https://arxiv.org/html/2606.18394#A3.T10)\.

##### Cumulative log\-probability scoring\.

accum\_logpachieves8\.15×8\.15\\timesspeedup \(average accepted lengthτ=9\.81\\tau=9\.81\) and is our production default\. It expands the tree by best\-first heap priority on cumulative draft log\-probability along each branch, prioritizing high\-likelihood continuations under the budget\.

##### Entropy\-only collapse\.

Entropy\-guided scoring \(frontier priority by per\-depth marginal entropy alone, with no cumulative\-log\-prob signal\) drops to4\.76×4\.76\\times\(−42%\-42\\%relative toaccum\_logp\), confirming that per\-depth entropy alone is insufficient to identify high\-acceptance branches\.

##### Hybrid scoring with increasingα\\alpha\.

Hybrid scoring uses∑ilog⁡ri\+α⋅Hi\\sum\_\{i\}\\log r\_\{i\}\+\\alpha\\cdot H\_\{i\}, with per\-depth entropyHiH\_\{i\}\. Atα≤1\\alpha\\leq 1, hybrid converges within∼1\.5%\{\\sim\}1\.5\\%ofaccum\_logp\(8\.158\.15–8\.27×8\.27\\times\); increasingα\\alphamonotonically degrades acceptance, reaching7\.42×7\.42\\timesatα=8\\alpha\{=\}8\. Theα→0\\alpha\\to 0limit reduces to pure cumulative log\-probability \(consistent withaccum\_logpsitting on the same plateau as low\-α\\alphahybrid\); theα→∞\\alpha\\to\\inftylimit corresponds to entropy\-only scoring \(consistent with the4\.76×4\.76\\timescollapse\)\.

Table 10:Tree\-construction algorithm ablation on MATH\-500 \(n=500n=500\) withJetFlowat the production setting \(causal head, LR3×10−43\{\\times\}10^\{\-4\}, Forward\-KL distillation,γ=0\\gamma=0\)\. Hybrid scoring is∑ilog⁡ri\+α⋅Hi\\sum\_\{i\}\\log r\_\{i\}\+\\alpha\\cdot H\_\{i\}with per\-depth entropyHiH\_\{i\}\. We report speedup and average accepted lengthτ\\tau\.

## Appendix DTraining Details

For both DFlash andJetFlow, we perform a learning\-rate sweep from1×10−41\{\\times\}10^\{\-4\}to1×10−31\{\\times\}10^\{\-3\}with five settings\. We find that3×10−43\{\\times\}10^\{\-4\}and6×10−46\{\\times\}10^\{\-4\}generally perform best, with different tasks favoring different choices\. Since their overall performance is comparable, we report results with3×10−43\{\\times\}10^\{\-4\}by default\. All draft\-head training runs are conducted on 8 H100 GPUs with a micro batch of 2\.

For a fair comparison with DFlash, theJetFlowdraft head is trained with block size 16, corresponding to a maximum tree depth of 16, and uses fused hidden features extracted from the frozen target model\. During training, each block keeps the first token as the anchor and replaces the remaining positions with mask\-token embeddings\. The causal head predicts all masked future tokens in parallel, while each position can attend only to the prefix and earlier positions within the same block, ensuring that the draft distribution follows the autoregressive order\. We sample random anchor positions from each sequence, use up to 512 anchors per example\. An example of the attention mask with 3 anchors for 3 blocks is provided in Figure[5](https://arxiv.org/html/2606.18394#A4.F5)\.

![Refer to caption](https://arxiv.org/html/2606.18394v1/figures/training_mask.png)Figure 5:Causal attention mask used for training with multiple sampled blocks\. Each query can attend to the full verified prefix and to the anchor plus earlier positions within its own block, but cannot attend to future positions or positions from other sampled blocks\.JetFlowreuses intermediate representations from the frozen target model as draft\-head context\. For Qwen3\-8B, we extract hidden states from target layers\{1,9,17,25,33\}\\\{1,9,17,25,33\\\}out of the 36\-layer target model, concatenate them along the channel dimension, and project the resulting5​d5dfeature back to hidden sized=4096d=4096through a bias\-free linear layer followed by RMSNorm\. The draft head is implemented as a lightweight Qwen3\-style decoder with 5 layers, 32 attention heads, 8 KV heads, head dimension 128, and MLP intermediate size 12288\. In each draft layer, the projected target feature is injected as contextual key/value states and concatenated with the draft\-token hidden states, allowing the causal draft head to condition on rich target\-model features while keeping the target model frozen\.

For ablations, we consider DFlash\-style exponential depth weighting\. Since depth weighting can implicitly bias the bidirectional DFlash head toward left\-to\-right prediction, in main results reported in Table[1](https://arxiv.org/html/2606.18394#S3.T1)and Table[2](https://arxiv.org/html/2606.18394#S3.T2), we remove it when isolating the effect of causal masking versus the native block\-diffusion head design\. For distillation runs, teacher logits are obtained from the frozen target model and aligned with the student prediction positions; the draft head is trained with a temperature\-scaled soft\-label distillation loss\.

![Refer to caption](https://arxiv.org/html/2606.18394v1/figures/blockwise_supervision.png)Figure 6:Each sampled block includes an anchor position and multiple future token positions\. The anchor is retained as block context and excluded from the loss, while loss is applied only to future token positions within each block\.
## Appendix EvLLM Implementation Details

We implement causal parallel tree drafting inside vLLM\. The draft head uses causal attention, while its candidates are organized as a speculative tree rather than a linear block\.

Given the target hidden states at the current decoding step, the proposer produces draft logits for a block of future positions\. For each position, it extracts the top\-kktokens and log\-probabilities, wherekkis the tree width, and constructs a speculative tree under a fixed node budget\. Our implementation supports several expansion scores, including cumulative log\-probability, entropy\-guided scoring, and hybrid scoring, as discussed in Appendix[C](https://arxiv.org/html/2606.18394#A3)\. We use cumulative log\-probability in our experiments\. The resulting tree stores token ids, parent indices, and node depths, with the root corresponding to the already\-sampled next token\.

During verification, the tree metadata is passed to the target model as parent indices and depths for all speculative nodes\. The target verifies all tree nodes in one forward pass using a tree attention mask\. To reduce verification overhead, we implement an optimized fused paged tree\-attention kernel: it builds the ancestor relation for speculative nodes, applies the tree mask inside attention, and verifies all candidates without materializing a dense per\-request mask\. Acceptance is computed by comparing each node token with the target model’s greedy prediction at its parent, then selecting the deepest fully accepted root\-to\-node path\. The accepted path is committed, and the target prediction at the final accepted node is used as the correction token\.

Table[11](https://arxiv.org/html/2606.18394#A5.T11)reports an example end\-to\-end vLLM run on HumanEval with a single H100 GPU\. The results show how tree budget interacts with serving batch size: larger budgets improve low\-batch latency by reducing verification rounds, while their relative gain decreases at larger batch sizes as verification and memory pressure become more prominent\.

Table 11:vLLM serving performance ofJetFlowon Math\-500 with Qwen3\-8B on a single H100 GPU, evaluated across batch sizes and tree budgets \(in parentheses\)\. Each setting reports end\-to\-end throughput over AR decoding in tokens per second \(TPS\)\.
## Appendix FAblation Study Details

We provide detailed ablation results for the training and generalization settings summarized in Section[3](https://arxiv.org/html/2606.18394#S3)\. Unless otherwise specified, ablations are conducted on the Nemotron Post\-Training Dataset V2 math split, using Qwen3\-8B with block size 16, node budget 255, accumulated draft log\-probability for tree construction, and the same evaluation protocol as the main experiments\.

##### Learning Rate\.

Table[3](https://arxiv.org/html/2606.18394#S3.T3)studies the effect of draft\-head learning rate\. Very small learning rates underfit the drafter, while performance plateaus around3×10−43\\times 10^\{\-4\}; larger rates remain competitive but do not provide consistent additional gains\.

##### Training Objective\.

Table[4](https://arxiv.org/html/2606.18394#S3.T4)compares SFT, forward\-KL distillation, and reverse\-KL distillation\. SFT and forward\-KL perform similarly, while reverse\-KL substantially degrades performance, suggesting that mode\-seeking distillation is poorly matched to budgeted tree drafting, where preserving multiple plausible continuations is important\.

##### Model Generalizability\.

Table[5](https://arxiv.org/html/2606.18394#S3.T5)evaluates whetherJetFlowgeneralizes beyond dense Qwen3\-8B\. On the Qwen3\-30B\-A3B MoE target model,JetFlowcontinues to outperform DFlash\-T under the same training recipe, showing that causal parallel tree drafting is not specific to a single model architecture\.

##### Training Data\.

Table[6](https://arxiv.org/html/2606.18394#S3.T6)compares regenerated target\-model sequences with direct corpus training\. Regenerated data gives the strongest performance by better matching the target model’s own generation distribution\. Corpus\-trainedJetFlowis weaker but still effective\.

## Appendix GEmpirical Per\-token Drafting Cost on Modern Hardware

Figure[2](https://arxiv.org/html/2606.18394#S2.F2)shows that speculative decoding scales better with draft length when the per\-token drafting cost is low\. In this appendix, we measure this cost for a practical DFlash\-style draft head on a single H200 NVL GPU\. We useNNto denote the draft depth, i\.e\., the number of draft tokens proposed in one block\. We sweep context length and draft depth with request batch size fixed to one, and report the per\-draft\-token cost ratio corresponding to the coefficientccin Eq\. equation[2](https://arxiv.org/html/2606.18394#S2.E2)\.

##### Setup\.

We profile a DFlash configuration using Qwen3\-8B as the target model andz\-lab/Qwen3\-8B\-DFlash\-b16as the draft head\. Each cell reports steady\-state latency after 10 warmup steps and 50 measurement steps, measured with CUDA events using the PyTorch SDPA backend\. For context lengthLLand draft depthNN, letTdraft​\(N,L\)T\_\{\\mathrm\{draft\}\}\(N,L\)denote the latency of one parallel draft\-head forward pass that proposesNNtokens, and letTverify​\(N,L\)T\_\{\\mathrm\{verify\}\}\(N,L\)denote the latency of one parallel target\-model verification pass over the sameNNcandidates\. Following Eq\. equation[2](https://arxiv.org/html/2606.18394#S2.E2), we define the per\-draft\-token cost coefficient as

c​\(N,L\)=Tdraft​\(N,L\)/NTverify​\(N,L\)=Tdraft​\(N,L\)N​Tverify​\(N,L\)\.c\(N,L\)=\\frac\{T\_\{\\mathrm\{draft\}\}\(N,L\)/N\}\{T\_\{\\mathrm\{verify\}\}\(N,L\)\}=\\frac\{T\_\{\\mathrm\{draft\}\}\(N,L\)\}\{N\\,T\_\{\\mathrm\{verify\}\}\(N,L\)\}\.The factor1/N1/Namortizes the parallel draft\-head forward cost across theNNproposed tokens\. Lowerccindicates cheaper marginal drafting relative to one target verification pass\.

##### Results\.

Table[12](https://arxiv.org/html/2606.18394#A7.T12)reportsccin percent over a sweep of context lengthsL∈\{128,…,4096\}L\\\!\\in\\\!\\\{128,\\dots,4096\\\}and draft depthsN∈\{1,2,4,8,16,32,64,128,256,512\}N\\\!\\in\\\!\\\{1,2,4,8,16,32,64,128,256,512\\\}\.

Table 12:Per\-draft\-token drafting cost ratioc=Tdraft/\(N​Tverify\)c=T\_\{\\mathrm\{draft\}\}/\(N\\,T\_\{\\mathrm\{verify\}\}\)\(%\) on a single H200 NVL GPU, sweeping context lengthLLand draft depthNN\. This is the cost coefficient used in Eq\. equation[2](https://arxiv.org/html/2606.18394#S2.E2)and Fig\.[2](https://arxiv.org/html/2606.18394#S2.F2)\. Lower is better\.Across the practically relevant regime \(L≤2048L\\\!\\leq\\\!2048,N≥16N\\\!\\geq\\\!16\), the per\-draft\-token cost is below1%1\\%of one target verification pass and decreases roughly as1/N1/N, since the parallel draft\-head cost is amortized across more proposed tokens\. AtL=1024L\\\!=\\\!1024, for instance,ccdrops from0\.845%0\.845\\%atN=16N\\\!=\\\!16to0\.054%0\.054\\%atN=256N\\\!=\\\!256, close to the ultra\-low\-cost regime in Fig\.[2](https://arxiv.org/html/2606.18394#S2.F2)\. Longer contexts increase the cost mildly, butccremains below0\.1%0\.1\\%forN≥256N\\\!\\geq\\\!256\.

##### Implication\.

These measurements support the main scaling argument: increasing the draft depthNNamortizes the draft\-head cost and lowers the effective per\-token cost coefficientcc\. On the same hardware,ccdecreases from≈6\.7%\\approx 6\.7\\%atN=2N\\\!=\\\!2to≈0\.05%\\approx 0\.05\\%atN=256N\\\!=\\\!256, moving from the typical\-cost regime toward the ultra\-low\-cost regime in Fig\.[2](https://arxiv.org/html/2606.18394#S2.F2)\. However, as Fig\.[2](https://arxiv.org/html/2606.18394#S2.F2)illustrates, low cost enables long\-draft scaling only when acceptance remains sufficiently high\. This motivates causal tree drafting as a way to improve acceptance quality in the low\-cost speculative decoding\.