PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding

# PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding
Source: [https://arxiv.org/html/2607.22634](https://arxiv.org/html/2607.22634)
Zheng Wang1,∗Zhifan Ye2,3,∗Qi Cheng3Yonggan Fu3Ziyan Wang2Feng Zhu2 Haozhe Zhao1Jan Kautz3Pavlo Molchanov3Humphrey Shi2,3Minjia Zhang1 1University of Illinois Urbana\-Champaign2Georgia Institute of Technology3NVIDIA ∗Equal contribution

###### Abstract

Abstract:Diffusion Large Language Models \(dLLMs\) have recently emerged as a promising alternative to autoregressive \(AR\) LLMs, offering parallel token generation\. Recent works have shown that dLLMs are particularly well\-suited as draft models for speculative decoding \(SD\), as they can efficiently generate an entire block of draft tokens in parallel within a single forward pass\. However, existing diffusion\-based drafting methods primarily rely on linear drafting \(single\-path\), despite diffusion models simultaneously producing multiple candidate tokens at multiple positions that naturally induce a large combinatorial space of possible decoding paths\. As a result, these methods can only explore a tiny fraction of the available candidate space, inherently limiting the achievable acceptance length and decoding efficiency\. To fully utilize this rich multi\-candidate structure, we apply tree\-based drafting to diffusion drafters, enabling the exploration of diverse candidate paths\. Nevertheless, we find that applying naive tree\-based drafting is suboptimal due to a fundamental mismatch between diffusion draft confidence and prefix\-based AR verification: diffusion marginals are inherently prefix\-blind, which can lead to unreliable path ranking\. We propose PRESTO, a principled framework that extends tree\-based drafting to diffusion drafters while resolving the fundamental mismatch between diffusion draft confidence and prefix\-based AR verification through*PREfix\-alignedScoring*and*priority\-basedTree search*for diffusion speculative decOding\. The key principles behind PRESTO are that \(1\) candidate ranking should align with the prefix\-based nature of AR verification, and \(2\) tree construction should prioritize candidate paths with high verification potential to maximize acceptance length\. We also demonstrate that PRESTO is designed to be a general tree drafting framework applicable to both dedicated diffusion drafter SD and self\-speculative dLLMs\. Extensive experiments show that PRESTO achieves an average of1\.5×\\timesend\-to\-end throughput speedup on the state\-of\-the\-art dedicated diffusion drafter SD and an average of1\.12×\\timeson self\-speculative diffusion LLMs across diverse benchmarks\.

![Refer to caption](https://arxiv.org/html/2607.22634v1/x1.png)Figure 1:Acceptance length comparison of PRESTO on diffusion\-based SD, including dedicated diffusion drafter SD \(Qwen3\-8B\-DFlash\) and self\-speculative dLLM \(Nemotron\-Labs\-Diffusion\-8B; marginal\-only scoring with no prefix\-signal\) across diverse benchmarks\. PRESTO consistently improves acceptance length over the baseline across all tasks, enabling more tokens to be accepted in each verification round\.## 1Introduction

Diffusion Large Language Models \(dLLMs\) have recently emerged as a promising alternative to autoregressive \(AR\) LLMs, introducing a new paradigm for multi\-token generation\[[1](https://arxiv.org/html/2607.22634#bib.bib1),[2](https://arxiv.org/html/2607.22634#bib.bib2),[3](https://arxiv.org/html/2607.22634#bib.bib3),[4](https://arxiv.org/html/2607.22634#bib.bib4),[5](https://arxiv.org/html/2607.22634#bib.bib5),[6](https://arxiv.org/html/2607.22634#bib.bib6)\]\. However, achieving high\-quality generation in practice typically requires iterative denoising procedures and carefully designed token unmasking strategies, which introduce additional computation and partially offset the efficiency gains\[[3](https://arxiv.org/html/2607.22634#bib.bib3),[7](https://arxiv.org/html/2607.22634#bib.bib7)\]\. Interestingly, recent works have shown that diffusion models are particularly well\-suited as draft models in speculative decoding\[[8](https://arxiv.org/html/2607.22634#bib.bib8),[9](https://arxiv.org/html/2607.22634#bib.bib9),[10](https://arxiv.org/html/2607.22634#bib.bib10),[11](https://arxiv.org/html/2607.22634#bib.bib11)\]\. Their ability to generate multiple candidate tokens in parallel naturally aligns with the proposal stage of speculative decoding \(SD\), enabling highly efficient draft token generation\. This insight has led to a growing line of work that leverages diffusion models as effective drafters, including dedicated diffusion\-based drafter for AR LLMs \(e\.g\., dFlash\[[10](https://arxiv.org/html/2607.22634#bib.bib10)\]\) and self\-speculative dLLMs \(e\.g\., TiDAR\[[8](https://arxiv.org/html/2607.22634#bib.bib8)\], I\-DLM\[[11](https://arxiv.org/html/2607.22634#bib.bib11)\], Nemotron\-Labs\-Diffusion\[[12](https://arxiv.org/html/2607.22634#bib.bib12)\]\)\. For example, dFlash demonstrates that well trained lightweight diffusion drafters can substantially accelerate speculative decoding through efficient parallel draft token generation while maintaining high draft quality, highlighting the practical potential of diffusion\-based drafting\.

However, existing diffusion\-based drafting methods primarily rely on*linear \(single\-path\) drafting*, where tokens are generated along a single trajectory and verified sequentially\. While this design is effective in AR settings, it is inherently suboptimal for diffusion models\. Due to the non\-autoregressive nature of diffusion generation, multiple plausible candidates can exist simultaneously, while the decoding path with the highest acceptance length may not correspond to the top\-1 marginal trajectory\. Empirically, we observe that even strong diffusion drafters in dFlash yield only modest accepted lengths under single\-path drafting \(See Figure[2](https://arxiv.org/html/2607.22634#S3.F2)\(a\)\), indicating that a significant portion of the proposal space remains unexplored\.

A natural direction to address this limitation is to move beyond linear drafting and explore multiple candidate paths jointly\. In AR settings, prior work has shown that introducing tree\-based drafting can substantially increase acceptance length by expanding multiple candidate paths in parallel, thereby improving the utilization of each verification step\[[13](https://arxiv.org/html/2607.22634#bib.bib13),[14](https://arxiv.org/html/2607.22634#bib.bib14),[15](https://arxiv.org/html/2607.22634#bib.bib15),[16](https://arxiv.org/html/2607.22634#bib.bib16),[17](https://arxiv.org/html/2607.22634#bib.bib17),[18](https://arxiv.org/html/2607.22634#bib.bib18)\]\. In particular, since the target model is often not fully compute\-bound during verification\[[8](https://arxiv.org/html/2607.22634#bib.bib8)\]for small batch size cases, allocating additional drafting compute to explore a richer candidate set can lead to more tokens being accepted per step\. The success of tree\-based drafting in AR settings raises a key research question:

*Can we leverage tree\-based drafting in diffusion drafters to further improve the end\-to\-end throughput of diffusion\-based speculative decoding by increasing acceptance length?*

While this direction is promising, we find that naively applying tree\-based drafting to diffusion models is often suboptimal\. Diffusion drafters generate position\-wise token distributions in parallel, without strict prefix dependencies, which enables efficient and flexible construction of tree\-structured candidates\. However, this same property also implies that the resulting scores are marginal token probabilities, rather than prefix\-conditioned likelihoods\. Consequently, these scores cannot accurately reflect how earlier token choices influence downstream acceptance \(e\.g\., a token that is likely in isolation may become unlikely once conditioned on the chosen prefix\), creating a fundamental mismatch between draft scoring and prefix\-based AR verification\.

In this work, we propose PRESTO, a principled tree\-based diffusion drafting framework via*prefix\-aligned scoring*and*priority\-based tree search*\. The key principles behind PRESTO are that \(1\) candidate ranking should align with the prefix\-based nature of AR verification, and \(2\) tree construction should prioritize candidate paths with high verification potential to maximize acceptance length\. Specifically, we minimally adjust the diffusion marginal by incorporating a prefix\-dependent correction, so that the resulting score better reflects the acceptance behavior under sequential verification\. This score is then used to guide a priority\-based expansion strategy, which allocates computation to high\-quality candidate paths during tree construction to maximize acceptance length\. As a result, PRESTO enables effective tree\-based diffusion drafting and significantly improves acceptance length and throughput on both dedicated diffusion\-based drafters for AR LLMs and self\-speculative dLLMs\. Specifically, we summarize our contributions as follows:

- •We identify a fundamental mismatch between diffusion\-based draft scoring and prefix\-based AR verification\.We show that while diffusion probabilities provide well\-calibrated marginal acceptance signals, they do not fully capture the prefix\-conditioned nature of AR verification, which can lead to suboptimal path ranking during tree construction and ultimately limit acceptance length\.
- •We design prefix\-aligned scoring to address the above mismatch\.We formulate a principled, minimal correction to diffusion marginals by incorporating prefix\-dependent signals, yielding a scoring function that is both faithful to the diffusion drafters and compatible with prefix\-based verification\.
- •We formalize and enable tree\-based speculative drafting for diffusion drafters\.Building upon the proposed prefix\-aligned scoring, we introduce PRESTO, a principled tree\-based diffusion drafting framework that efficiently explores high\-quality draft paths through priority\-based tree expansion, substantially improving acceptance length and decoding throughput\.
- •PRESTO generalizes across both dedicated and self\-speculative diffusion SD frameworks\.We further extend PRESTO to self\-speculative dLLMs as a drop\-in replacement for linear drafting in the parallelized verify\-and\-draft pipeline, making PRESTO the first tree\-based drafting framework applicable to both dedicated diffusion drafters and self\-speculative diffusion LLMs\. Across diverse benchmarks and model sizes, PRESTO achieves an average of1\.5×\\timesthroughput speedup on dFlash and an average of1\.12×\\timesthroughput speedup on Nemotron\-Labs\-Diffusion over linear drafting baselines\.

## 2Preliminaries

### 2\.1Tree\-based Speculative Decoding

Speculative decoding accelerates autoregressive generation of a target LLMpTp\_\{T\}by leveraging a lightweight draft model\. To improve acceptance, the draft model can construct a token tree encoding multiple candidate sequences\[[13](https://arxiv.org/html/2607.22634#bib.bib13),[19](https://arxiv.org/html/2607.22634#bib.bib19),[18](https://arxiv.org/html/2607.22634#bib.bib18)\]\. The target LLM verifies candidates in a prefix\-based manner: starting from the root, it accepts tokens sequentially until a mismatch occurs\.

Acceptance factorization\.For a candidate pathP=\(x1,…,xk\)P=\(x\_\{1\},\\ldots,x\_\{k\}\), letai∈\{0,1\}a\_\{i\}\\in\\\{0,1\\\}denote the indicator that theii\-th token is accepted\. Acceptance proceeds left\-to\-right, so by the chain rule

Pr⁡\(P​accepted\)=∏i=1kPr⁡\(ai=1\|a<i=1,x≤i\),\\Pr\(P\\text\{ accepted\}\)=\\prod\_\{i=1\}^\{k\}\\Pr\\\!\\bigl\(a\_\{i\}=1\\,\\bigm\|\\,a\_\{<i\}=1,\\,x\_\{\\leq i\}\\bigr\),\(1\)wherea<i=1a\_\{<i\}=1denotes the event that all earlier tokens have been accepted\. Each factor depends on the realized prefixx<ix\_\{<i\}, making verification intrinsically prefix\-conditional\.

Tree construction objective\.LetℬB=\{𝒯:\|𝒯\|≤B\}\\mathcal\{B\}\_\{B\}=\\\{\\mathcal\{T\}:\|\\mathcal\{T\}\|\\leq B\\\}denote token trees with at mostBBnodes where𝒯\\mathcal\{T\}is the tree budget, and letα𝒯​\(P\)\\alpha\_\{\\mathcal\{T\}\}\(P\)denote the longest prefix ofPPcontained in𝒯\\mathcal\{T\}\. The ideal objective is

𝒯∗∈arg⁡max𝒯∈ℬB⁡𝔼P∼pT​\[α𝒯​\(P\)\]\.\\mathcal\{T\}^\{\*\}\\in\\arg\\max\_\{\\mathcal\{T\}\\in\\mathcal\{B\}\_\{B\}\}\\mathbb\{E\}\_\{P\\sim p\_\{T\}\}\[\\alpha\_\{\\mathcal\{T\}\}\(P\)\]\.\(2\)Direct optimization is infeasible, as evaluatingpTp\_\{T\}requires target\-model forward passes\. We use a surrogate distributionp~\\tilde\{p\}over candidate paths\. A standard decomposition \(Appendix[A](https://arxiv.org/html/2607.22634#A1)\) yields

𝔼P∼p~​\[α𝒯​\(P\)\]=∑u∈𝒯p~​\(u\),\\mathbb\{E\}\_\{P\\sim\\tilde\{p\}\}\[\\alpha\_\{\\mathcal\{T\}\}\(P\)\]=\\sum\_\{u\\in\\mathcal\{T\}\}\\tilde\{p\}\(u\),\(3\)wherep~​\(u\)\\tilde\{p\}\(u\)denotes the probability that prefixuumatches the sampled path\. This reduces tree construction to selecting prefixes with large probability mass\. Sincep~​\(u\)\\tilde\{p\}\(u\)factorizes along the prefix,

log⁡p~​\(u\)=∑i=1\|u\|log⁡p~​\(ui∣u<i\),\\log\\tilde\{p\}\(u\)=\\sum\_\{i=1\}^\{\|u\|\}\\log\\tilde\{p\}\(u\_\{i\}\\mid u\_\{<i\}\),\(4\)providing an additive scoring rule that enables priority\-based expansion\.

### 2\.2Prefix\-Aligned Surrogate Distributions

###### Definition 2\.1\(Prefix\-Aligned Surrogate\)\.

A surrogate distributionp~\\tilde\{p\}over𝒱k\\mathcal\{V\}^\{k\}is*prefix\-aligned*if there exists some positioni≥2i\\geq 2such that the conditional distributionp~​\(ui∣u<i\)\\tilde\{p\}\(u\_\{i\}\\mid u\_\{<i\}\)is not invariant to the prefix, i\.e\., there existu<i≠u<i′u\_\{<i\}\\neq u^\{\\prime\}\_\{<i\}with

p~​\(ui∣u<i\)≠p~​\(ui∣u<i′\)\.\\tilde\{p\}\(u\_\{i\}\\mid u\_\{<i\}\)\\neq\\tilde\{p\}\(u\_\{i\}\\mid u^\{\\prime\}\_\{<i\}\)\.\(5\)

The nontriviality clause is essential\. Under an independence assumption, any joint distribution admits a trivial prefix\-blind factorization of the formp~​\(u\)=∏ip~i​\(ui\)\\tilde\{p\}\(u\)=\\prod\_\{i\}\\tilde\{p\}\_\{i\}\(u\_\{i\}\), where each token is scored independently of the prefix\. Such decompositions assign identical per\-position contributions regardless of the realized prefix and therefore fail to capture the prefix\-conditioned structure in Eq\. \([1](https://arxiv.org/html/2607.22634#S2.E1)\)\.

### 2\.3Autoregressive vs\. Diffusion Surrogate Distributions

Existing tree drafting methods for AR drafters implicitly assume autoregressive prefix\-conditioned scores\[[13](https://arxiv.org/html/2607.22634#bib.bib13),[15](https://arxiv.org/html/2607.22634#bib.bib15),[19](https://arxiv.org/html/2607.22634#bib.bib19),[18](https://arxiv.org/html/2607.22634#bib.bib18)\]\. This assumption fundamentally breaks under diffusion drafting\.

Autoregressive drafting yields a prefix\-aligned surrogate\.An AR draft modelpDp\_\{D\}defines a prefix\-conditioned distribution over sequences via the chain rule:

p~AR​\(P\):=pD​\(P\)=∏i=1kpD​\(xi∣x<i\)\.\\tilde\{p\}\_\{\\text\{AR\}\}\(P\):=p\_\{D\}\(P\)=\\prod\_\{i=1\}^\{k\}p\_\{D\}\(x\_\{i\}\\mid x\_\{<i\}\)\.\(6\)Each conditionalpD​\(xi∣x<i\)p\_\{D\}\(x\_\{i\}\\mid x\_\{<i\}\)is a genuine function of the prefix, a well\-trained AR drafter approximates the target’s prefix conditionals, sop~AR\\tilde\{p\}\_\{\\text\{AR\}\}is prefix\-aligned \(Definition[2\.1](https://arxiv.org/html/2607.22634#S2.Thmdefinition1)\) and its log\-probability factorization in Eq\. \([6](https://arxiv.org/html/2607.22634#S2.E6)\) provides an effective surrogate score for tree construction\.

Diffusion drafting yields a prefix\-blind marginal surrogate\.A diffusion drafter produces a collection of*position\-wise marginals*q1,…,qkq\_\{1\},\\ldots,q\_\{k\}over𝒱\\mathcal\{V\}\. Eachqi​\(xi\)q\_\{i\}\(x\_\{i\}\)is the marginal distribution of positionii; it does not condition on the prefixx<i=\(x1,…,xi−1\)x\_\{<i\}=\(x\_\{1\},\\ldots,x\_\{i\-1\}\)\.111The dependence on the prompt and previously generated context, encoded in a shared latent state, is suppressed in this section for notational clarity\.The induced factorized distribution over candidate paths is

p~diff​\(P\):=∏i=1kqi​\(xi\)\.\\tilde\{p\}\_\{\\text\{diff\}\}\(P\):=\\prod\_\{i=1\}^\{k\}q\_\{i\}\(x\_\{i\}\)\.\(7\)
The marginal surrogatep~diff\\tilde\{p\}\_\{\\text\{diff\}\}is not prefix\-aligned: each conditionalp~diff​\(xi∣x<i\)=qi​\(xi\)\\tilde\{p\}\_\{\\text\{diff\}\}\(x\_\{i\}\\mid x\_\{<i\}\)=q\_\{i\}\(x\_\{i\}\)takes the same value regardless ofx<ix\_\{<i\}, violating the nontriviality clause of Definition[2\.1](https://arxiv.org/html/2607.22634#S2.Thmdefinition1)\. The structural mismatch with prefix\-based verification is*intrinsic*to the marginal factorization, not an artifact of any particular drafter or training procedure\. In Section[3](https://arxiv.org/html/2607.22634#S3), we will show that while this mismatch does not preclude diffusion\-based drafting from being effective, it can affect the quality of path ranking\.

## 3Observations

![Refer to caption](https://arxiv.org/html/2607.22634v1/x2.png)Figure 2:\(a\) Linear drafting leaves room for improvement via multi\-path exploration\. \(b\) The predicted confidence fromqdq\_\{d\}closely matches the empirical acceptance rate\. \(c\) Within fixedqdq\_\{d\}bins, empirical acceptance rates increase from low to highρd\\rho\_\{d\}terciles, indicating that prefix\-conditioned compatibility captures predictive signal beyond the marginal confidenceqdq\_\{d\}\. \(d\)ρd\\rho\_\{d\}provides a correction direction aligned with the error ofqdq\_\{d\}, although it remains noisy \(correlation<1<1\)\.In this section, we use Qwen3\-4B\-DFlash with mask block size being 16 as an example to analyze the characteristics of diffusion\-based draft scoring\. To study the impact of tree construction, we adopt the tree expansion strategy from EAGLE\[[15](https://arxiv.org/html/2607.22634#bib.bib15),[16](https://arxiv.org/html/2607.22634#bib.bib16)\]as the default tree construction algorithm\. For each candidate token, we record its diffusion draft probabilityqdq\_\{d\}, prefix\-conditioned scoreρd\\rho\_\{d\}, and final acceptance outcome\. The prefix\-conditioned scoreρd\\rho\_\{d\}is instantiated using annn\-gram model to capture compatibility with prefix\-based AR verification\. Unless otherwise specified, results are aggregated across datasets\.

Observation 1: Linear drafting leaves substantial room for improvement through multi\-path exploration\.We first compare the acceptance length achieved by linear drafting with an oracle upper bound equivalent to exhaustive multi\-path exploration\. In our setting, the diffusion drafter generates a block of1616draft positions simultaneously, naturally exposing a rich multi\-candidate draft space within a single forward pass\. To estimate the achievable upper bound, at each position, we retain the top\-10 candidates from the diffusion marginal distribution and perform a position\-wise oracle matching analysis\. Specifically, we check whether the ground\-truth token at each position appears in the corresponding top\-10 candidate set, and then compute the longest contiguous prefix length covered by these matches\.222The reported upper bound is equivalent to the result that would be obtained by exhaustive enumeration\. However, we do not explicitly enumerate all101610^\{16\}paths\. Instead, we exploit the fact that acceptance is determined by whether the ground\-truth token is covered by the top\-k candidate set at each position\. This allows us to compute the same oracle upper bound efficiently through position\-wise matching\.As shown in Figure[2](https://arxiv.org/html/2607.22634#S3.F2)\(a\), a substantial gap consistently exists between linear drafting and this upper bound across all tasks\. For example, on GSM8K the acceptance length improves from6\.56\.5to9\.789\.78, while on Math500 it increases from7\.97\.9to11\.1611\.16\. Similar trends are observed on AIME25, HumanEval, and MBPP\. These results suggest that existing linear drafting strategies fail to fully exploit the rich combinatorial candidate space already produced by diffusion drafters, leaving significant room for improvement via multi\-path exploration\.

Observation 2: Marginal confidence correlates with overall acceptance\.We further examine whether the diffusion drafter provides informative estimates of downstream verification success\. Specifically, conditioned on the preceding prefix being accepted by the verifier, we compute the diffusion marginal confidenceqdq\_\{d\}for the next draft token using the exponentiated marginal probability under the diffusion distribution, and compare it against the empirical acceptance rate observed during verification\. As shown in Figure[2](https://arxiv.org/html/2607.22634#S3.F2)\(b\), the acceptance rate increases monotonically withqdq\_\{d\}, demonstrating a strong positive correlation between diffusion confidence and actual verification success\. This suggests that the diffusion marginal distribution already captures meaningful global acceptance trends, and that higher\-confidence draft tokens are substantially more likely to be accepted\. At the same time, the calibration curve consistently remains below the idealy=xy=xline, indicating that marginal confidence alone is still imperfect and does not fully characterize verifier acceptance\. Nevertheless,qdq\_\{d\}provides a strong global signal for estimating acceptance likelihood and serves as a useful foundation for candidate ranking\.

Observation 3: Marginal signals are insufficient but can be improved by prefix conditioning\.Despite its strong global correlation with acceptance,qdq\_\{d\}alone is insufficient for reliable token\-level ranking: as shown in Figure[2](https://arxiv.org/html/2607.22634#S3.F2)\(c\), tokens with similarqdq\_\{d\}can yield substantially different acceptance outcomes when grouped byρd\\rho\_\{d\}\. This implies thatqdq\_\{d\}is not a sufficient statistic for predicting acceptance, and thatρd\\rho\_\{d\}provides complementary, prefix\-conditioned information\. This additional prefix\-conditioned signal also remains directionally informative: Figure[2](https://arxiv.org/html/2607.22634#S3.F2)\(d\) shows that the correction termδ=log⁡ρd−log⁡qd\\delta=\\log\\rho\_\{d\}\-\\log q\_\{d\}positively correlates with the true residualrq=log⁡pT−log⁡qdr\_\{q\}=\\log p\_\{T\}\-\\log q\_\{d\}\. This indicates thatρd\\rho\_\{d\}provides a correction direction that aligns with the true error of the draft model\. Importantly, the correlation remains significantly below 1, indicating thatρd\\rho\_\{d\}is not a direct proxy for the target distribution\. Instead, it acts as a noisy but directionally aligned correction signal\.

Insights\.The draft probabilityqdq\_\{d\}provides informative estimates of global acceptance, but is still insufficient for accurate path ranking\. The prefix\-conditioned signalρd\\rho\_\{d\}offers complementary information that aligns with the error ofqdq\_\{d\}, but is itself noisy\. Together, these observations suggest that*an effective scoring should preserveqdq\_\{d\}while incorporating a controlled correction based onρd\\rho\_\{d\}*\.

## 4PRESTO

We now introduce PRESTO, a tree\-based diffusion drafting framework that augments diffusion drafter with prefix\-aligned information, enabling more effective candidate exploration under prefix\-conditioned AR verification\. In Section[4\.1](https://arxiv.org/html/2607.22634#S4.SS1), we introduce prefix\-aligned scoring that addresses the aforementioned structural mismatch\. In Section[4\.2](https://arxiv.org/html/2607.22634#S4.SS2), we describe tree drafting via priority\-based tree search\. In Section[4\.3](https://arxiv.org/html/2607.22634#S4.SS3), we elaborate how to extend tree drafting to self\-speculative dLLMs\.

### 4\.1Prefix\-Aligned Scoring via Minimal Correction

From surrogate distribution to scoring function\.Section[2](https://arxiv.org/html/2607.22634#S2)reduces tree construction to ranking prefixes bylog⁡p~​\(u\)\\log\\tilde\{p\}\(u\)\(Eq\. \([4](https://arxiv.org/html/2607.22634#S2.E4)\)\)\. For diffusion drafting, the natural marginal surrogatep~diff​\(P\)=∏iqi​\(xi\)\\tilde\{p\}\_\{\\text\{diff\}\}\(P\)=\\prod\_\{i\}q\_\{i\}\(x\_\{i\}\)is prefix\-blind \(Eq\. \([7](https://arxiv.org/html/2607.22634#S2.E7)\)\), creating a mismatch with prefix\-based verification\. We therefore seek a prefix\-aligned surrogatep~⋆\\tilde\{p\}^\{\\star\}that preserves the calibrated drafter signal while incorporating a prefix\-conditioned correction in terms of the observations and insights from Section[3](https://arxiv.org/html/2607.22634#S3)\.

Prefix\-aligned surrogate\.Letqd​\(t\)q\_\{d\}\(t\)denote the diffusion drafter’s marginal at positiondd, and letρd​\(t∣cd\)\\rho\_\{d\}\(t\\mid c\_\{d\}\)denote a tractable prefix\-conditioned signal, wherecd=\(t1,…,td−1\)c\_\{d\}=\(t\_\{1\},\\ldots,t\_\{d\-1\}\)is the within\-block prefix\. We define the prefix\-aligned conditional as a multiplicative combination:

pd⋆​\(t∣cd\)∝qd​\(t\)​ρd​\(t∣cd\)λd,p^\{\\star\}\_\{d\}\(t\\mid c\_\{d\}\)\\propto q\_\{d\}\(t\)\\,\\rho\_\{d\}\(t\\mid c\_\{d\}\)^\{\\lambda\_\{d\}\},\(8\)whereλd≥0\\lambda\_\{d\}\\geq 0controls the strength of the prefix correction\. The induced joint surrogate distribution is

p~⋆​\(P\):=∏d=1kpd⋆​\(td∣cd\),\\tilde\{p\}^\{\\star\}\(P\):=\\prod\_\{d=1\}^\{k\}p^\{\\star\}\_\{d\}\(t\_\{d\}\\mid c\_\{d\}\),\(9\)which is prefix\-aligned \(Definition[2\.1](https://arxiv.org/html/2607.22634#S2.Thmdefinition1)\) wheneverλd\>0\\lambda\_\{d\}\>0andρd\\rho\_\{d\}depends nontrivially oncdc\_\{d\}\. Eq\. \([9](https://arxiv.org/html/2607.22634#S4.E9)\) thereby provides a candidate surrogate that addresses the structural mismatch ofp~diff\\tilde\{p\}\_\{\\text\{diff\}\}\. This form can also be interpreted as the solution to a KL\-regularized objective that balances closeness toqdq\_\{d\}with alignment toρd\\rho\_\{d\}\(see Appendix[B](https://arxiv.org/html/2607.22634#A2)\)\.

For simplicity, we use the unnormalized log\-score induced by Eq\. \([8](https://arxiv.org/html/2607.22634#S4.E8)\)\. Specifically, we define the token\-level score assd,t​\(cd\)=log⁡qd​\(t\)\+λd​log⁡ρd​\(t∣cd\)s\_\{d,t\}\(c\_\{d\}\)=\\log q\_\{d\}\(t\)\+\\lambda\_\{d\}\\log\\rho\_\{d\}\(t\\mid c\_\{d\}\), and the path score

S​\(P\)=∑d=1ksd,td​\(cd\),cd=\(t1,…,td−1\)\.S\(P\)=\\sum\_\{d=1\}^\{k\}s\_\{d,t\_\{d\}\}\(c\_\{d\}\),\\qquad c\_\{d\}=\(t\_\{1\},\\ldots,t\_\{d\-1\}\)\.\(10\)Eq\. \([10](https://arxiv.org/html/2607.22634#S4.E10)\) decomposes additively over depth, supporting incremental priority\-based tree expansion\.

### 4\.2Priority\-Based Tree Construction

Algorithm 1Beam Search with Global Retention1:procedureBeamSearch\(

s,B,b,W,Ds,B,b,W,D\)

2:

𝒫←\{root\},𝒯←\{root\}\\mathcal\{P\}\\leftarrow\\\{\\text\{root\}\\\},\\ \\mathcal\{T\}\\leftarrow\\\{\\text\{root\}\\\}
3:for

d=1d=1to

DDwhile

\|𝒯\|<B\|\\mathcal\{T\}\|<Bdo

4:

𝒫new←∅\\mathcal\{P\}\_\{\\text\{new\}\}\\leftarrow\\emptyset
5:foreach

v∈𝒫v\\in\\mathcal\{P\}do

6:

𝒞←Expand​\(v,b,s\)\\mathcal\{C\}\\leftarrow\\textsc\{Expand\}\(v,b,s\)
7:

𝒫new←𝒫new∪𝒞\\mathcal\{P\}\_\{\\text\{new\}\}\\leftarrow\\mathcal\{P\}\_\{\\text\{new\}\}\\cup\\mathcal\{C\}
8:

𝒯←𝒯∪𝒞\\mathcal\{T\}\\leftarrow\\mathcal\{T\}\\cup\\mathcal\{C\}
9:endfor

10:

𝒫←TopW​\(𝒫new,W\)\\mathcal\{P\}\\leftarrow\\textsc\{TopW\}\(\\mathcal\{P\}\_\{\\text\{new\}\},W\)
11:endfor

12:return

TopB​\(𝒯,B\)\\textsc\{TopB\}\(\\mathcal\{T\},B\)
13:endprocedure

Algorithm 2Best\-First Search1:procedureBestFirst\(

s,B,bs,B,b\)

2:

𝒫←\{root\},𝒯←\{root\}\\mathcal\{P\}\\leftarrow\\\{\\text\{root\}\\\},\\ \\mathcal\{T\}\\leftarrow\\\{\\text\{root\}\\\}
3:

S​\(root\)←0S\(\\text\{root\}\)\\leftarrow 0
4:while

\|𝒯\|<B\|\\mathcal\{T\}\|<Bdo

5:

v⋆←arg⁡maxu∈𝒫⁡S​\(u\)v^\{\\star\}\\leftarrow\\arg\\max\_\{u\\in\\mathcal\{P\}\}S\(u\)
6:

𝒫←𝒫∖\{v⋆\}\\mathcal\{P\}\\leftarrow\\mathcal\{P\}\\setminus\\\{v^\{\\star\}\\\}
7:

𝒞←Expand​\(v⋆,b,s\)\\mathcal\{C\}\\leftarrow\\textsc\{Expand\}\(v^\{\\star\},b,s\)
8:Insert

𝒞\\mathcal\{C\}into priority queue

𝒫\\mathcal\{P\}
9:

𝒯←𝒯∪𝒞\\mathcal\{T\}\\leftarrow\\mathcal\{T\}\\cup\\mathcal\{C\}
10:endwhile

11:⊳\\trianglerighttree construction complete

12:return

TopB​\(𝒯,B\)\\textsc\{TopB\}\(\\mathcal\{T\},B\)
13:endprocedure

From the surrogate objective to priority\-based expansion\.Section[2](https://arxiv.org/html/2607.22634#S2)shows that maximizing the surrogate objective in Eq\. \([3](https://arxiv.org/html/2607.22634#S2.E3)\) reduces to selecting the top\-BBprefixes by their log\-probabilities underp~⋆\\tilde\{p\}^\{\\star\}\. Equivalently, the optimal tree consists of theBBprefixesuuwith the largestlog⁡p~⋆​\(u\)=S​\(u\)\\log\\tilde\{p\}^\{\\star\}\(u\)=S\(u\), whereSSis the path score in Eq\. \([10](https://arxiv.org/html/2607.22634#S4.E10)\)\. However, exhaustively expanding and ranking all candidate continuations is computationally infeasible due to the exponentially growing search space and large vocabulary branching factor\. Instead, since the path score decomposes additively over depth, high\-scoring partial prefixes are more likely to remain among the top\-ranked candidates after further expansion\. This naturally motivates a greedy priority\-based expansion strategy for diffusion drafting: maintain a frontier of expandable nodes and iteratively expand the highest\-scoring node\.

We identify each tree nodevvwith the unique path from the root tovv, denotedP​\(v\)=\(t1,…,td​\(v\)\)P\(v\)=\(t\_\{1\},\\ldots,t\_\{d\(v\)\}\), whered​\(v\):=\|P​\(v\)\|d\(v\):=\|P\(v\)\|is the depth ofvv\. The cumulative path score isS​\(v\):=S​\(P​\(v\)\)=∑i=1d​\(v\)si,ti​\(ci\)S\(v\):=S\(P\(v\)\)=\\sum\_\{i=1\}^\{d\(v\)\}s\_\{i,t\_\{i\}\}\(c\_\{i\}\)withci=\(t1,…,ti−1\)c\_\{i\}=\(t\_\{1\},\\ldots,t\_\{i\-1\}\)\. The root hasd​\(root\)=0,S​\(root\)=0d\(\\text\{root\}\)=0,S\(\\text\{root\}\)=0\. At each iteration, the algorithm:

1. 1\.Selects a frontier nodev⋆v^\{\\star\}according to a search policyπ\\pi;
2. 2\.Expandsv⋆v^\{\\star\}by adding its top\-bbchildren, ranked by the token scoresd​\(v⋆\)\+1,t​\(cd​\(v⋆\)\+1\)s\_\{d\(v^\{\\star\}\)\+1,t\}\(c\_\{d\(v^\{\\star\}\)\+1\}\)at depthd​\(v⋆\)\+1d\(v^\{\\star\}\)\+1;
3. 3\.Updates the frontier with the newly added children, removingv⋆v^\{\\star\}from the expandable set\.

The iteration terminates when\|𝒯\|=B\|\\mathcal\{T\}\|=B\. BecauseSSis additive over depth, each child’s cumulative score is computed incrementally from its parent’s, avoiding redundant computation\.

We consider two widely used instantiations of priority\-based expansion under the path scoreSS: beam search with global retention \(Algorithm[1](https://arxiv.org/html/2607.22634#alg1)\) and best\-first search \(BFS\) \(Algorithm[2](https://arxiv.org/html/2607.22634#alg2)\)\. Both prioritize high\-SSpartial paths but differ in how the expansion budget is allocated across depth: BFS maintains a global priority over the entire frontier, while beam search restricts expansion to the top\-WWin\-beam nodes per depth and retains out\-of\-beam nodes for the final top\-BBselection\. Beam search can thus be interpreted as a width\-constrained approximation to BFS\.

Table 1:Acceptance length comparison\.#### BFS vs\. Beam Search\.

Table[1](https://arxiv.org/html/2607.22634#S4.T1)reports acceptance length under the two policies across budget sizes\. At small to moderate budgets \(B≤512B\\leq 512\), best\-first and beam search yield similar acceptance length, as beam search’s per\-depth budget is sufficient to cover most high\-priority expansions that best\-first search would select under global prioritization\. AtB=1024B=1024, best\-first outperforms beam search by a small margin: with a larger budget, the preferred expansion pattern can become less uniform across depth, which a global priority handles more flexibly than a fixed per\-depth beam\. In practice, verification shifts from memory\-bound to compute\-bound onceBBexceeds a hardware\-dependent threshold, after which forward\-pass cost grows linearly withBBand erodes the speedup\. We therefore operate atB≤512B\\leq 512, which sits within the memory\-bound regime where the two policies are comparable, and adopt beam search as the default search strategy\.

![Refer to caption](https://arxiv.org/html/2607.22634v1/x3.png)Figure 3:Application of PRESTO to hybrid dLLMs under self\-speculative decoding with a tree budget of 4 drafting tokens\. The committed prefix reuses the exact KV cache for AR verification of the previous draft tree, while trailing mask\-token slots are simultaneously denoised for diffusion\-based speculation within the same forward pass\. Speculative outputs attached to rejected branches are discarded, while valid speculation slots along the accepted path are used for next\-round tree construction\. The resulting draft tree is flattened into draft tokens for the next forward pass under AR tree attention\. PRESTO is the first framework to explore tree\-based drafting for self\-speculative dLLMs\. The detailed process is provided in Appendix[C](https://arxiv.org/html/2607.22634#A3)\.

### 4\.3Extending Tree Construction to hybrid dLLMs

The proposed tree\-construction framework can be naturally extended to support quadratic self\-speculative decoding in hybrid dLLMs\[[8](https://arxiv.org/html/2607.22634#bib.bib8),[12](https://arxiv.org/html/2607.22634#bib.bib12)\]333Nemotron\-Labs\-Diffusion primarily highlights a linear self\-speculation decoding mode, where the model first generates draft tokens through diffusion decoding and then verifies them autoregressively using the same model\. Applying PRESTO to this setting is straightforward and identical to the dedicated diffusion drafter setting discussed earlier, as PRESTO only modifies the diffusion drafting process while leaving the verification procedure unchanged\. In addition to linear self\-speculation, the Nemotron\-Labs\-Diffusion paper also introduces a quadratic self\-speculation variant\[[12](https://arxiv.org/html/2607.22634#bib.bib12)\]which is further explored in this work\., as illustrated in Figure[3](https://arxiv.org/html/2607.22634#S4.F3)\. Quadratic self\-speculation makes autoregressive verification and diffusion\-based drafting parallelized within a single model forward pass: previously drafted tokens are verified autoregressively, while new draft candidates are simultaneously generated at trailing mask positions through diffusion decoding\.

PRESTO preserves this original verify\-and\-draft pipeline without introducing additional model forwards\. During each forward pass, the committed prefix reuses the exact KV cache for AR verification of the previous draft tree, while the trailing mask\-token slots attached to speculative branches are simultaneously denoised for next\-round drafting\. After verification, only the mask slots attached to the longest accepted path remain valid speculation positions, while the slots attached to rejected branches are discarded since their predictions were conditioned on prefixes that are no longer part of the committed context\. PRESTO then extracts top\-kkdiffusion candidates from the valid mask slots and organizes them into a new draft tree using the proposed prefix\-aligned scoring and priority\-based tree expansion\. The resulting tree is flattened into draft tokens and verified in the next forward pass using the AR tree attention mechanism\.

This extension is particularly well\-suited for self\-speculative dLLMs because in standard linear self\-speculative decoding, speculative tokens attached to rejected branches are discarded after verification, leading to substantial wasted computation\. PRESTO alleviates this inefficiency by organizing diffusion outputs into a tree structure, allowing diverse speculative candidates generated within the same forward pass to be more effectively utilized during verification\. As shown in Section[5\.2](https://arxiv.org/html/2607.22634#S5.SS2), replacing linear drafting with PRESTO consistently improves acceptance length and decoding throughput across different datasets and model sizes\.

## 5Empirical Validation

In this section, we demonstrate that PRESTO significantly improves the efficiency of diffusion\-based speculative decoding across both standard and self\-speculative diffusion based SD\.

### 5\.1Experiment Setup

Models and Tasks\.We evaluate PRESTO on two settings: \(i\) standard speculative decoding with a separate diffusion drafter and AR target: dFlash with Qwen3\-4B, Qwen3\-8B, and Qwen3\-Coder\-30B\-A3B backbones; and \(ii\) hybrid dLLM: Nemotron\-Labs\-Diffusion\-8B at both linear self\-speculation and quadratic self\-speculation modes\. Evaluation spans three task categories:mathematical reasoning\(GSM8K\[[20](https://arxiv.org/html/2607.22634#bib.bib20)\], MATH\[[21](https://arxiv.org/html/2607.22634#bib.bib21)\], AIME24\[[22](https://arxiv.org/html/2607.22634#bib.bib22)\], AIME25\[[23](https://arxiv.org/html/2607.22634#bib.bib23)\]\),code\(HumanEval\[[24](https://arxiv.org/html/2607.22634#bib.bib24)\], MBPP\[[25](https://arxiv.org/html/2607.22634#bib.bib25)\], LiveCodeBench\[[26](https://arxiv.org/html/2607.22634#bib.bib26)\]\), andconversation\(MT\-Bench\[[27](https://arxiv.org/html/2607.22634#bib.bib27)\], Alpaca\[[28](https://arxiv.org/html/2607.22634#bib.bib28)\]\)\. We report average acceptance length \(τ\\tau\) and throughput \(token/s\) speedup over the AR baseline\. All experiments run on NVIDIA B200 GPUs\.

Implementations\.We implement PRESTO using two runtime backends\. For dFlash and Nemotron\-Labs\-Diffusion under quadratic self\-speculation, we use a PyTorch implementation with FlexAttention\[[29](https://arxiv.org/html/2607.22634#bib.bib29)\]for target model verification with tree\-structured attention mask\[[13](https://arxiv.org/html/2607.22634#bib.bib13),[14](https://arxiv.org/html/2607.22634#bib.bib14),[15](https://arxiv.org/html/2607.22634#bib.bib15)\]\. For Nemotron\-Labs\-Diffusion under linear self\-speculation, we integrate PRESTO into the official SGLang implementation with FlashInfer attention\[[30](https://arxiv.org/html/2607.22634#bib.bib30)\]\. We instantiate the prefix\-aligned score using a lightweight n\-gram model withn=3n=3444Following the n\-gram model construction in\[[31](https://arxiv.org/html/2607.22634#bib.bib31)\]\.\. Unless otherwise specified, we setλd=0\.2\\lambda\_\{d\}=0\.2, batch size to 1, and use beam search with global retention, using beam width 10 and top\-kkexpansion \(k=10k=10\) per position\. We evaluate PRESTO under both deterministic \(T=0T=0\) and stochastic \(T=1T=1\) decoding\. For dFlash, we sweep tree budgets𝒯∈\{128,256,512\}\\mathcal\{T\}\\in\\\{128,256,512\\\}and report the best\-performing configuration in terms of end\-to\-end throughput\. For Nemotron\-Labs\-Diffusion, we use a fixed tree budget𝒯\\mathcal\{T\}of 32 following its speculative block size\.

### 5\.2Experiment Results

Table 2:Decoding speedup over baseline and average acceptance length \(τ\\tau\) on Qwen3 models with thinking mode disabled and a maximum of 2048 generated tokens\.PRESTO on dFlash\.As illustrated in Table[2](https://arxiv.org/html/2607.22634#S5.T2), PRESTO consistently achieves substantial improvements in the average acceptance lengthτ\\tauacross all models and tasks, which directly translates into higher end\-to\-end throughput speedup\. For example, on Qwen3\-4B withT=0T=0,τ\\tauincreases from7\.97\.9to10\.910\.9on MATH\-500 and from7\.47\.4to10\.310\.3on AIME25, while on code benchmarks it improves from6\.96\.9to10\.010\.0on LCB and from6\.76\.7to9\.99\.9on HumanEval, leading to corresponding speedup gains across all settings\. The largest improvements are observed on math and code benchmarks, whereτ\\tautypically increases by\+2\+2to\+3\+3tokens\. These trends remain stable across model scales and decoding regimes: even for stronger drafters such as Qwen3\-8B and under stochastic decoding \(T=1T=1\), PRESTO continues to deliver comparable relative gains\. These results indicate that PRESTO effectively pushes the acceptance length of dFlash closer to its upper bound, delivering consistent and substantial end\-to\-end throughput speedup gains\. The end\-to\-end throughput results based on torch implementation is shown in Appendix[D](https://arxiv.org/html/2607.22634#A4)\.

Table 3:Decoding speedup over baseline, average acceptance length \(τ\\tau\), and absolute throughput on Nemotron\-Labs\-Diffusion\-8B at linear self speculation decoding mode with a maximum of 2048 generated tokens\.Table 4:Decoding speedup over baseline, average acceptance length \(τ\\tau\), and absolute throughput on Nemotron\-Labs\-Diffusion\-8B at quadratic self speculation decoding mode with a maximum of 2048 generated tokens\.PRESTO on Nemotron\-Labs\-Diffusion \(NLD\)\.As shown in Table[3](https://arxiv.org/html/2607.22634#S5.T3), PRESTO consistently improves both throughput and average acceptance length over the official linear self\-speculative decoding implementation of Nemotron\-Labs\-Diffusion across all benchmarks and decoding temperatures\. By replacing linear drafting with tree\-based exploration while preserving the original verification pipeline, PRESTO remains effective across a wide range of tasks\. Under deterministic decoding \(T=0T=0\), PRESTO increases average end\-to\-end throughput from4\.6×4\.6\\timesto4\.9×4\.9\\timeswhile improving average acceptance length from8\.88\.8to9\.99\.9, with consistent gains observed on all benchmarks\. Under stochastic decoding \(T=1T=1\), where a single drafted trajectory is more likely to diverge from the target model, the benefits become more pronounced: average end\-to\-end throughput improves from2\.1×2\.1\\timesto2\.5×2\.5\\times, while average acceptance length increases from4\.24\.2to5\.15\.1\. Across individual benchmarks, PRESTO achieves up to1\.29×1\.29\\timesend\-to\-end throughput improvement and increases acceptance length by as much as39%39\\%on Alpaca under stochastic decoding \(T=1\)\. These results demonstrate that the benefits of tree\-based drafting persist even within a production\-grade, highly optimized inference engine\. The end\-to\-end throughput results based on SGLang is shown in Appendix[E](https://arxiv.org/html/2607.22634#A5)\.

PRESTO on Nemotron\-Labs\-Diffusion \(NLD\) at quadratic self speculation mode\.We further evaluate PRESTO under the quadratic self\-speculative decoding setting introduced in Nemotron\-Labs\-Diffusion\. As Table[4](https://arxiv.org/html/2607.22634#S5.T4)shows, across all benchmarks and decoding settings, PRESTO also consistently outperforms vanilla quadratic self\-speculative decoding in both end\-to\-end throughput and average acceptance length\. Under deterministic decoding \(T=0T=0\), PRESTO provides stable gains across math, code, and chat tasks, improving average speedup from2\.6×2\.6\\timesto3\.5×3\.5\\timeson NLD\-8B, while also consistently increasing acceptance length from4\.74\.7to5\.65\.6\. The largest gains are observed on math benchmarks, where PRESTO achieves up to2\.04×2\.04\\timeson GSM8k for NLD\-8B\. Gains are largest under stochastic decoding \(T=1T=1\), where single\-path drafting degrades sharply from trajectory mismatch and accumulated uncertainty, while PRESTO ’s tree\-structured exploration retains multiple plausible continuations\. On NLD\-8B, this raises average end\-to\-end speedup from1\.4×1\.4\\timesto3\.0×3\.0\\timesand nearly doubles acceptance length from2\.42\.4to4\.74\.7, reaching up to2\.9×2\.9\\timesthroughput and over2×2\\timesacceptance length over vanilla self\-speculative decoding\. These positive results indicate the significant benefits of applying PRESTO to quadratic self\-speculation\.

We note that the throughput for dFlash and NLD quadratic self speculation can be further improved through system\-level optimizations, including custom kernels, optimized KV\-cache management, and efficient request scheduling\. The goal here is not to fully optimize absolute throughput, but rather to provide a first\-hand comparison between linear and tree\-based drafting for diffusion drafters under the same native PyTorch implementation\.

## 6Discussions and Ablations

![Refer to caption](https://arxiv.org/html/2607.22634v1/x4.png)

![Refer to caption](https://arxiv.org/html/2607.22634v1/x5.png)

![Refer to caption](https://arxiv.org/html/2607.22634v1/x6.png)

Figure 4:Top: Effectiveness of prefix\-conditioned compatibility signals\. Incorporating the prefix\-conditioned termρd\\rho\_\{d\}consistently improves acceptance length across datasets and tree token budgets\.Middle: Sensitivity analysis ofλd\\lambda\_\{d\}under different datasets and tree token budgets\. Moderate values ofλd\\lambda\_\{d\}consistently achieve the best performance, withλd≈0\.2\\lambda\_\{d\}\\approx 0\.2providing the most stable gains across settings\.Bottom left: Per\-step latency breakdown across different model scales\. Target\-model verification dominates the overall decoding cost, while tree\-related operations introduce only negligible overhead\.Bottom right: Distribution of acceptance lengths on GSM8K\. PRESTO significantly shifts the distribution toward longer accepted prefixes and substantially increases the frequency of near\-complete and full\-block acceptance\. All results are obtained using Qwen3\-4B\-DFlash with diffusion block size 16\.Effectiveness of prefix\-conditioned signal\.We first ablate the contribution of the prefix\-conditioned compatibility termρd\\rho\_\{d\}by comparing our full scoring function against a variant that only relies on the diffusion marginal confidenceqdq\_\{d\}\. As shown in Figure[4](https://arxiv.org/html/2607.22634#S6.F4)top, incorporatingρd\\rho\_\{d\}consistently improves the average acceptance length across all tasks and tree token budgets\. The gains are particularly noticeable on reasoning and coding benchmarks such as GSM8K, Math500, and HumanEval, where prefix consistency plays a more important role in determining downstream verification success\. These results suggest that diffusion marginal confidence alone is insufficient for reliably ranking candidate paths, while the prefix\-conditioned signal provides complementary information that better aligns tree expansion with the autoregressive verification process\.

Sensitivity ofλd\\lambda\_\{d\}\.Figure[4](https://arxiv.org/html/2607.22634#S6.F4)middle shows the impact of varyingλd\\lambda\_\{d\}under different datasets and tree token budgets\. Across nearly all settings, moderate values ofλd\\lambda\_\{d\}consistently achieve the best acceptance length, withλd≈0\.2\\lambda\_\{d\}\\approx 0\.2providing the most stable performance\. Whenλd\\lambda\_\{d\}is too small, the prefix\-conditioned signal is underutilized and the scoring function degenerates toward pure diffusion confidence ranking\. In contrast, overly largeλd\\lambda\_\{d\}overemphasizesρd\\rho\_\{d\}and suppresses the diffusion marginal confidenceqdq\_\{d\}, leading to unstable candidate ranking and substantial degradation in acceptance length\. Importantly, the optimal range remains highly consistent across different tasks and tree sizes, suggesting that the proposed scoring function is robust\.

Entropy Adaptiveλd\\lambda\_\{d\}vs\. Fixedλd\\lambda\_\{d\}\.We further explore an entropy\-adaptiveλd\\lambda\_\{d\}that dynamically balancesqdq\_\{d\}andρd\\rho\_\{d\}according to the uncertainty of the diffusion drafter: confident predictions rely more onqdq\_\{d\}, while uncertain predictions assign larger weights toρd\\rho\_\{d\}\(details in Appendix[F](https://arxiv.org/html/2607.22634#A6)\)\. As shown in Figure[4](https://arxiv.org/html/2607.22634#S6.F4)top, entropy\-adaptiveλd\\lambda\_\{d\}provides only marginal improvements over a fixedλd\\lambda\_\{d\}, and the overall trends remain highly similar across datasets\. This suggests that although uncertainty\-aware weighting is intuitively appealing, a simple fixedλd\\lambda\_\{d\}is already sufficient to capture most of the benefits of prefix\-conditioned ranking in practice, while avoiding additional complexity\.

Overhead Analysis\.As shown in Figure[4](https://arxiv.org/html/2607.22634#S6.F4)bottom left, target\-model verification overwhelmingly dominates the overall decoding latency across all model scales, accounting for 81%–97% of the total runtime\. In comparison, tree\-related operations, including tree construction, candidate filtering, and attention\-mask preparation, contribute only 0\.7%–4% of the per\-step latency\. This demonstrates that the proposed tree construction introduces negligible system overhead relative to the target\-model forward pass\. Moreover, the overhead becomes proportionally smaller for larger models, indicating favorable scalability of the framework\.

Acceptance Length Distribution Analysis\.Figure[4](https://arxiv.org/html/2607.22634#S6.F4)bottom right further illustrates how PRESTO reshapes the acceptance\-length distribution\. Compared with vanilla diffusion drafting, PRESTO substantially reduces short\-acceptance cases while significantly increasing the frequency of long\-prefix and full\-block acceptance\. In particular, the probability of achieving near\-complete or full\-block acceptance increases dramatically, indicating that the proposed prefix\-conditioned tree construction is more effective at identifying globally consistent candidate paths rather than locally confident but incompatible drafts\. This distributional shift directly explains the throughput improvements observed in the main results\.

## 7Related Work

### 7\.1Speculative Decoding

Speculative decoding\[[32](https://arxiv.org/html/2607.22634#bib.bib32),[33](https://arxiv.org/html/2607.22634#bib.bib33)\]accelerates AR models by drafting tokens with a fast proposer and verifying them in parallel with the target model, provably preserving the target distribution\. This guarantee relies on the target model possessing a well\-trained verify distribution\. A line of follow\-up work eliminates the external drafter or improves draft quality: Medusa\[[14](https://arxiv.org/html/2607.22634#bib.bib14)\]augments the base LLM with multiple prediction heads and uses tree attention for parallel verification; the EAGLE family\[[34](https://arxiv.org/html/2607.22634#bib.bib34),[15](https://arxiv.org/html/2607.22634#bib.bib15),[16](https://arxiv.org/html/2607.22634#bib.bib16)\]exploits feature\-level context from the frozen target model, with EAGLE\-1 predicting future hidden states to boost acceptance, EAGLE\-2 introducing adaptive drafting trees, and EAGLE\-3 refining training objectives to scale speedups; and SpecInfer\[[13](https://arxiv.org/html/2607.22634#bib.bib13)\]proposes tree\-structured verification\. Multi\-token prediction \(MTP\)\[[35](https://arxiv.org/html/2607.22634#bib.bib35)\]trains models to predict multiple future tokens simultaneously\. Despite these advances, most existing methods rely on autoregressive drafting, which remains inherently sequential and limits attainable speedups\.

### 7\.2Tree Construction in Speculative Decoding

A complementary line of work focuses on how draft tokens are organized for parallel verification\. Medusa\[[14](https://arxiv.org/html/2607.22634#bib.bib14)\]popularized tree\-structured drafting by packing multiple candidate continuations into a single token tree and verifying them in one forward pass via a topology\-aware attention mask\. However, this tree mask is hand\-crafted and static\. SpecInfer\[[13](https://arxiv.org/html/2607.22634#bib.bib13)\]introduced token\-tree verification with provable distribution preservation\. Subsequent work moves beyond fixed shapes: Sequoia\[[18](https://arxiv.org/html/2607.22634#bib.bib18)\]formulates tree construction as a dynamic\-programming problem, jointly optimizing the tree topology and a hardware\-aware tree size for a given accelerator\. SpecExec\[[17](https://arxiv.org/html/2607.22634#bib.bib17)\]pushes draft trees to massively parallel sizes \(thousands of nodes\) to amortize the cost of offloaded weights on consumer hardware based on Best First Search like tree construction algorihtm\. EAGLE\-2\[[15](https://arxiv.org/html/2607.22634#bib.bib15)\]observes that draft acceptance is context\-dependent and uses the calibrated confidence scores of the EAGLE drafter to construct a dynamic draft tree per step via beam search like tree expension\. OPT\-Tree\[[36](https://arxiv.org/html/2607.22634#bib.bib36)\]formalizes the objective as maximizing the expected acceptance length and searches for the adaptive tree structure that attains it under a node budget\. These methods consistently show that careful tree drafting contributes a substantial fraction of end\-to\-end speedup\. However, they all assume an autoregressive drafter that produces tokens sequentially with conditional probabilities along each branch, and extending tree construction to diffusion\-based drafters, whose draft tokens are produced in parallel remains an open problem that our work directly addresses\.

### 7\.3Diffusion\-based Speculative Decoding

Recent work explores using diffusion models as drafters within speculative decoding, combining the parallelism of diffusion drafting with the quality guarantee of AR verification\. TiDAR\[[8](https://arxiv.org/html/2607.22634#bib.bib8)\]jointly trains diffusion and autoregressive objectives in a sequence\-level hybrid, enabling parallel “thinking” via diffusion and sequential “talking” via autoregressive decoding, though final generation quality is not yet lossless\. DiffuSpec\[[9](https://arxiv.org/html/2607.22634#bib.bib9)\]and SpecDiff\-2\[[37](https://arxiv.org/html/2607.22634#bib.bib37)\]employ large pretrained dLLMs as speculative drafters, with inference\-time search or train–test alignment to improve acceptance\. However, these approaches primarily focus on improving the quality or alignment of the diffusion drafter itself, often relying on massive drafters \(e\.g\., 7B parameters\), which incur substantial memory and latency overhead\. While they achieve long acceptance lengths, the high drafting cost can offset the practical speedups in real\-world serving scenarios\. In contrast, PRESTO studies how to more effectively exploit the rich multi\-position candidate space produced by diffusion drafting through prefix\-aware tree construction and priority\-based expansion, making it complementary to stronger diffusion drafters and existing diffusion\-based speculative decoding systems\. DFlash\[[10](https://arxiv.org/html/2607.22634#bib.bib10)\]proposes to employ a lightweight block\-diffusion drafter conditioned on context features extracted from the target model, generating an entire block of draft tokens in a single forward pass and reporting over6×6\\timeslossless acceleration\. Similarly, DART\[[31](https://arxiv.org/html/2607.22634#bib.bib31)\]performs parallel logit prediction over multiple masked positions and assembles drafts via N\-gram\-guided tree pruning\. However, its tree construction relies on a fixed, hand\-tuned scoring rule with hard\-coded hyperparameters, which leaves no principled way to adapt the tree shape to different budgets, drafters, or target models without re\-tuning\. In contrast, PRESTO formulates tree construction as priority\-based expansion under a unified path score, yielding a single algorithmic framework that naturally generalizes across different tree budgets and diffusion drafter designs\.

## 8Conclusion

In this work, we propose PRESTO, a principled framework for tree\-based speculative decoding with diffusion language models\. We identify a fundamental mismatch between diffusion draft scoring and prefix\-based autoregressive verification: while diffusion probabilities provide strong marginal plausibility signals, they are inherently prefix\-blind and therefore insufficient for reliable path ranking in tree\-based drafting\. To address this issue, we introduce a prefix\-aligned scoring mechanism together with a priority\-based tree search strategy, enabling effective exploration of high\-quality candidate paths during diffusion drafting\.

## 9Limitations

PRESTO requires a tractable prefix\-aligned signalρd\\rho\_\{d\}to complement the diffusion drafter’s prefix\-blind marginals\. In this work, we employ a lightweight n\-gram model, which is empirically effective, but only captures short\-range lexical compatibility\. Exploring richer prefix\-conditioned signals while preserving low overhead remains an important future direction\. We hope this motivates further efforts to inject prefix\-aligned signals into diffusion marginals\.

We additionally explored an entropy\-adaptive variant ofλd\\lambda\_\{d\}, but observed only marginal improvements over a fixed coefficient, suggesting that more sophisticated adaptation strategies may require learned or target\-aware scheduling mechanisms\.

Finally, we primarily evaluate PRESTO under single\-request decoding settings using the native PyTorch implementation except NLD at linear self speculation decoding mode\. Extending PRESTO to production\-scale serving frameworks such as vLLM and SGLang, together with characterizing its behavior under larger batch sizes, longer contexts, and highly optimized serving systems, remains meaningful future work\.

## References

- \[1\]Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji\-Rong Wen, and Chongxuan Li\.Large language diffusion models, 2025\.
- \[2\]Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong\.Dream 7b: Diffusion large language models, 2025\.
- \[3\]Chengyue Wu, Hao Zhang, Shuchen Xue, Shizhe Diao, Yonggan Fu, Zhijian Liu, Pavlo Molchanov, Ping Luo, Song Han, and Enze Xie\.Fast\-dllm v2: Efficient block\-diffusion llm, 2025\.
- \[4\]Tiwei Bie, Maosong Cao, Xiang Cao, Bingsen Chen, Fuyuan Chen, Kun Chen, Lun Du, Daozhuo Feng, Haibo Feng, Mingliang Gong, Zhuocheng Gong, Yanmei Gu, Jian Guan, Kaiyuan Guan, Hongliang He, Zenan Huang, Juyong Jiang, Zhonghui Jiang, Zhenzhong Lan, Chengxi Li, Jianguo Li, Zehuan Li, Huabin Liu, Lin Liu, Guoshan Lu, Yuan Lu, Yuxin Ma, Xingyu Mou, Zhenxuan Pan, Kaida Qiu, Yuji Ren, Jianfeng Tan, Yiding Tian, Zian Wang, Lanning Wei, Tao Wu, Yipeng Xing, Wentao Ye, Liangyu Zha, Tianze Zhang, Xiaolu Zhang, Junbo Zhao, Da Zheng, Hao Zhong, Wanli Zhong, Jun Zhou, Junlin Zhou, Liwang Zhu, Muzhi Zhu, and Yihong Zhuang\.Llada2\.1: Speeding up text diffusion via token editing, 2026\.
- \[5\]Shuang Cheng, Yihan Bian, Dawei Liu, Linfeng Zhang, Qian Yao, Zhongbo Tian, Wenhai Wang, Qipeng Guo, Kai Chen, Biqing Qi, and Bowen Zhou\.Sdar: A synergistic diffusion\-autoregression paradigm for scalable sequence generation, 2025\.
- \[6\]Yonggan Fu, Lexington Whalen, Zhifan Ye, Xin Dong, Shizhe Diao, Jingyu Liu, Chengyue Wu, Hao Zhang, Enze Xie, Song Han, Maksim Khadkevich, Jan Kautz, Yingyan Celine Lin, and Pavlo Molchanov\.Efficient\-dlm: From autoregressive to diffusion language models, and beyond in speed, 2026\.
- \[7\]Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, and Enze Xie\.Fast\-dllm: Training\-free acceleration of diffusion llm by enabling kv cache and parallel decoding, 2025\.
- \[8\]Jingyu Liu, Xin Dong, Zhifan Ye, Rishabh Mehta, Yonggan Fu, Vartika Singh, Jan Kautz, Ce Zhang, and Pavlo Molchanov\.Tidar: Think in diffusion, talk in autoregression, 2025\.
- \[9\]Guanghao Li, Zhihui Fu, Min Fang, Qibin Zhao, Ming Tang, Chun Yuan, and Jun Wang\.Diffuspec: Unlocking diffusion language models for speculative decoding, 2025\.
- \[10\]Jian Chen, Yesheng Liang, and Zhijian Liu\.Dflash: Block diffusion for flash speculative decoding, 2026\.
- \[11\]Yifan Yu, Yuqing Jian, Junxiong Wang, Zhongzhu Zhou, Donglin Zhuang, Xinyu Fang, Sri Yanamandra, Xiaoxia Wu, Qingyang Wu, Shuaiwen Leon Song, Tri Dao, Ben Athiwaratkun, James Zou, Fan Lai, and Chenfeng Xu\.Introspective diffusion language models, 2026\.
- \[12\]Yonggan Fu, Lexington Whalen, Abhinav Garg, Chengyue Wu, Maksim Khadkevich, Nicolai Oswald, Enze Xie, Daniel Egert, Sharath Turuvekere Sreenivas, Shizhe Diao, Chenhan Yu, Ye Yu, Weijia Chen, Sajad Norouzi, Jingyu Liu, Shiyi Lan, Ligeng Zhu, Jin Wang, Jindong Jiang, Morteza Mardani, Mehran Maghoumi, Song Han, Ante Jukic, Nima Tajbakhsh, Jan Kautz, and Pavlo Molchanov\.Nemotron\-labs\-diffusion: A tri\-mode language model unifying autoregressive, diffusion, and self\-speculation decoding\.May 2026\.Technical report\.
- \[13\]Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Zhengxin Zhang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, Chunan Shi, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia\.Specinfer: Accelerating 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, Volume 3, ASPLOS ’24, page 932–949\. ACM, April 2024\.
- \[14\]Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D\. Lee, Deming Chen, and Tri Dao\.Medusa: Simple llm inference acceleration framework with multiple decoding heads, 2024\.
- \[15\]Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang\.Eagle\-2: Faster inference of language models with dynamic draft trees, 2024\.
- \[16\]Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang\.Eagle\-3: Scaling up inference acceleration of large language models via training\-time test, 2025\.
- \[17\]Ruslan Svirschevski, Avner May, Zhuoming Chen, Beidi Chen, Zhihao Jia, and Max Ryabinin\.Specexec: Massively parallel speculative decoding for interactive llm inference on consumer devices, 2024\.
- \[18\]Zhuoming Chen, Avner May, Ruslan Svirschevski, Yuhsun Huang, Max Ryabinin, Zhihao Jia, and Beidi Chen\.Sequoia: Scalable, robust, and hardware\-aware speculative decoding, 2025\.
- \[19\]Ziteng Sun, Ananda Theertha Suresh, Jae Hun Ro, Ahmad Beirami, Himanshu Jain, and Felix Yu\.Spectr: Fast speculative decoding via optimal transport, 2024\.
- \[20\]Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman\.Training verifiers to solve math word problems, 2021\.
- \[21\]Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe\.Let’s verify step by step, 2023\.
- \[22\]Art of Problem Solving\.Aime problems and solutions\.[https://artofproblemsolving\.com/wiki/index\.php/AIME](https://artofproblemsolving.com/wiki/index.php/AIME), 2024\.Accessed: 2025\-04\-20\.
- \[23\]MAA\.American invitational mathematics examination \(aime\)\.[https://maa\.org/math\-competitions/aime](https://maa.org/math-competitions/aime)\.Mathematics Competition Series; n\.d\.a\.
- \[24\]Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert\-Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N\. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba\.Evaluating large language models trained on code, 2021\.
- \[25\]Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, and Charles Sutton\.Program synthesis with large language models, 2021\.
- \[26\]Naman Jain, King Han, Alex Gu, Wen\-Ding Li, Fanjia Yan, Tianjun Zhang, Sida Wang, Armando Solar\-Lezama, Koushik Sen, and Ion Stoica\.Livecodebench: Holistic and contamination free evaluation of large language models for code, 2024\.
- \[27\]Lianmin Zheng, Wei\-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric P\. Xing, Hao Zhang, Joseph E\. Gonzalez, and Ion Stoica\.Judging llm\-as\-a\-judge with mt\-bench and chatbot arena, 2023\.
- \[28\]Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B\. Hashimoto\.Stanford alpaca: An instruction\-following llama model\.[https://github\.com/tatsu\-lab/stanford\_alpaca](https://github.com/tatsu-lab/stanford_alpaca), 2023\.
- \[29\]Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He\.Flex attention: A programming model for generating optimized attention kernels, 2024\.
- \[30\]Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, and Luis Ceze\.Flashinfer: Efficient and customizable attention engine for llm inference serving, 2025\.
- \[31\]Fuliang Liu, Xue Li, Ketai Zhao, Yinxi Gao, Ziyan Zhou, Zhonghui Zhang, Zhibin Wang, Wanchun Dou, Sheng Zhong, and Chen Tian\.Dart: Diffusion\-inspired speculative decoding for fast llm inference, 2026\.
- \[32\]Yaniv Leviathan, Matan Kalman, and Yossi Matias\.Fast inference from transformers via speculative decoding, 2023\.
- \[33\]Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean\-Baptiste Lespiau, Laurent Sifre, and John Jumper\.Accelerating large language model decoding with speculative sampling, 2023\.
- \[34\]Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang\.Eagle: Speculative sampling requires rethinking feature uncertainty, 2025\.
- \[35\]Mohammad Samragh, Arnav Kundu, David Harrison, Kumari Nishu, Devang Naik, Minsik Cho, and Mehrdad Farajtabar\.Your llm knows the future: Uncovering its multi\-token prediction potential, 2025\.
- \[36\]Jikai Wang, Yi Su, Juntao Li, Qingrong Xia, Zi Ye, Xinyu Duan, Zhefeng Wang, and Min Zhang\.Opt\-tree: Speculative decoding with adaptive draft tree structure, 2025\.
- \[37\]Jameson Sandler, Jacob K\. Christopher, Thomas Hartvigsen, and Ferdinando Fioretto\.Specdiff\-2: Scaling diffusion drafter alignment for faster speculative decoding, 2025\.

## Appendix APrefix\-Mass Decomposition of the Tree Objective

Let𝒱\\mathcal\{V\}denote the token vocabulary and let𝒯\\mathcal\{T\}be a prefix\-closed tree whose nodes are finite token sequences\. Define

𝒯\+:=𝒯∖\{∅\}\\mathcal\{T\}^\{\+\}:=\\mathcal\{T\}\\setminus\\\{\\emptyset\\\}as the set of non\-root nodes\. For each nodeu=\(u1,…,u\|u\|\)∈𝒯\+u=\(u\_\{1\},\\ldots,u\_\{\|u\|\}\)\\in\\mathcal\{T\}^\{\+\}, let\|u\|\|u\|denote its depth\.

For a sampled path

P=\(x1,…,xk\)∼p~,P=\(x\_\{1\},\\ldots,x\_\{k\}\)\\sim\\tilde\{p\},define the length\-ddprefix

P≤d:=\(x1,…,xd\)\.P\_\{\\leq d\}:=\(x\_\{1\},\\ldots,x\_\{d\}\)\.The accepted\-prefix length under tree𝒯\\mathcal\{T\}is

α𝒯​\(P\):=max⁡\{d:P≤d∈𝒯\}\.\\alpha\_\{\\mathcal\{T\}\}\(P\):=\\max\\\{d:P\_\{\\leq d\}\\in\\mathcal\{T\}\\\}\.
Since𝒯\\mathcal\{T\}is prefix\-closed, the eventP≤d∈𝒯P\_\{\\leq d\}\\in\\mathcal\{T\}impliesP≤j∈𝒯P\_\{\\leq j\}\\in\\mathcal\{T\}for allj≤dj\\leq d\. Hence,

α𝒯​\(P\)=∑d=1k𝟏​\{P≤d∈𝒯\}\.\\alpha\_\{\\mathcal\{T\}\}\(P\)=\\sum\_\{d=1\}^\{k\}\\mathbf\{1\}\\\{P\_\{\\leq d\}\\in\\mathcal\{T\}\\\}\.Equivalently,

α𝒯​\(P\)=∑u∈𝒯\+𝟏​\{P≤\|u\|=u\}\.\\alpha\_\{\\mathcal\{T\}\}\(P\)=\\sum\_\{u\\in\\mathcal\{T\}^\{\+\}\}\\mathbf\{1\}\\\{P\_\{\\leq\|u\|\}=u\\\}\.
Taking expectation with respect toP∼p~P\\sim\\tilde\{p\}and applying linearity of expectation,

𝔼P∼p~​\[α𝒯​\(P\)\]=∑u∈𝒯\+PrP∼p~⁡\(P≤\|u\|=u\)\.\\mathbb\{E\}\_\{P\\sim\\tilde\{p\}\}\[\\alpha\_\{\\mathcal\{T\}\}\(P\)\]=\\sum\_\{u\\in\\mathcal\{T\}^\{\+\}\}\\Pr\_\{P\\sim\\tilde\{p\}\}\(P\_\{\\leq\|u\|\}=u\)\.
Define the surrogate prefix mass

p~​\(u\):=PrP∼p~⁡\(P≤\|u\|=u\)\.\\tilde\{p\}\(u\):=\\Pr\_\{P\\sim\\tilde\{p\}\}\(P\_\{\\leq\|u\|\}=u\)\.Then

𝔼P∼p~​\[α𝒯​\(P\)\]=∑u∈𝒯\+p~​\(u\)\.\\mathbb\{E\}\_\{P\\sim\\tilde\{p\}\}\[\\alpha\_\{\\mathcal\{T\}\}\(P\)\]=\\sum\_\{u\\in\\mathcal\{T\}^\{\+\}\}\\tilde\{p\}\(u\)\.
If the root node is excluded by convention, then𝒯\+=𝒯\\mathcal\{T\}^\{\+\}=\\mathcal\{T\}, yielding Eq\. \([3](https://arxiv.org/html/2607.22634#S2.E3)\)\.

## Appendix BKL\-Regularized Derivation of the Prefix\-Aligned Surrogate

We consider the variational objective

minp∈Δ​\(𝒱\)⁡KL​\(p∥qd​\(⋅\)\)−λd​𝔼t∼p​\[log⁡ρd​\(t∣cd\)\],\\min\_\{p\\in\\Delta\(\\mathcal\{V\}\)\}\\;\\mathrm\{KL\}\\\!\\left\(p\\,\\\|\\,q\_\{d\}\(\\cdot\)\\right\)\-\\lambda\_\{d\}\\mathbb\{E\}\_\{t\\sim p\}\\bigl\[\\log\\rho\_\{d\}\(t\\mid c\_\{d\}\)\\bigr\],\(11\)whereqdq\_\{d\}denotes the diffusion drafter marginal,ρd\(⋅∣cd\)\\rho\_\{d\}\(\\cdot\\mid c\_\{d\}\)is a prefix\-conditioned compatibility signal, andΔ​\(𝒱\)\\Delta\(\\mathcal\{V\}\)denotes the probability simplex over the vocabulary\. We assumeqd​\(t\)\>0q\_\{d\}\(t\)\>0andρd​\(t∣cd\)\>0\\rho\_\{d\}\(t\\mid c\_\{d\}\)\>0on the candidate support\.

Expanding the KL divergence,

KL​\(p∥qd\)=∑tp​\(t\)​log⁡p​\(t\)qd​\(t\),\\mathrm\{KL\}\(p\\\|q\_\{d\}\)=\\sum\_\{t\}p\(t\)\\log\\frac\{p\(t\)\}\{q\_\{d\}\(t\)\},the objective becomes

ℒ​\(p\)=∑tp​\(t\)​log⁡p​\(t\)qd​\(t\)−λd​∑tp​\(t\)​log⁡ρd​\(t∣cd\)\.\\mathcal\{L\}\(p\)=\\sum\_\{t\}p\(t\)\\log\\frac\{p\(t\)\}\{q\_\{d\}\(t\)\}\-\\lambda\_\{d\}\\sum\_\{t\}p\(t\)\\log\\rho\_\{d\}\(t\\mid c\_\{d\}\)\.
Including the normalization constraint∑tp​\(t\)=1\\sum\_\{t\}p\(t\)=1with Lagrange multiplierμ\\mu, we obtain

𝒥​\(p\)=∑tp​\(t\)​log⁡p​\(t\)qd​\(t\)−λd​∑tp​\(t\)​log⁡ρd​\(t∣cd\)\+μ​\(∑tp​\(t\)−1\)\.\\mathcal\{J\}\(p\)=\\sum\_\{t\}p\(t\)\\log\\frac\{p\(t\)\}\{q\_\{d\}\(t\)\}\-\\lambda\_\{d\}\\sum\_\{t\}p\(t\)\\log\\rho\_\{d\}\(t\\mid c\_\{d\}\)\+\\mu\\left\(\\sum\_\{t\}p\(t\)\-1\\right\)\.
Taking derivatives with respect top​\(t\)p\(t\),

∂𝒥∂p​\(t\)=log⁡p​\(t\)−log⁡qd​\(t\)\+1−λd​log⁡ρd​\(t∣cd\)\+μ\.\\frac\{\\partial\\mathcal\{J\}\}\{\\partial p\(t\)\}=\\log p\(t\)\-\\log q\_\{d\}\(t\)\+1\-\\lambda\_\{d\}\\log\\rho\_\{d\}\(t\\mid c\_\{d\}\)\+\\mu\.
Setting the derivative to zero yields

log⁡pd⋆​\(t\)=log⁡qd​\(t\)\+λd​log⁡ρd​\(t∣cd\)\+C,\\log p^\{\\star\}\_\{d\}\(t\)=\\log q\_\{d\}\(t\)\+\\lambda\_\{d\}\\log\\rho\_\{d\}\(t\\mid c\_\{d\}\)\+C,whereCCabsorbs constants independent oftt\. Exponentiating both sides,

pd⋆​\(t\)∝qd​\(t\)​ρd​\(t∣cd\)λd\.p^\{\\star\}\_\{d\}\(t\)\\propto q\_\{d\}\(t\)\\,\\rho\_\{d\}\(t\\mid c\_\{d\}\)^\{\\lambda\_\{d\}\}\.Equivalently,

pd⋆​\(t∣cd\)=qd​\(t\)​ρd​\(t∣cd\)λdZd​\(cd\),Zd​\(cd\)=∑t′∈𝒱qd​\(t′\)​ρd​\(t′∣cd\)λd\.p^\{\\star\}\_\{d\}\(t\\mid c\_\{d\}\)=\\frac\{q\_\{d\}\(t\)\\rho\_\{d\}\(t\\mid c\_\{d\}\)^\{\\lambda\_\{d\}\}\}\{Z\_\{d\}\(c\_\{d\}\)\},\\qquad Z\_\{d\}\(c\_\{d\}\)=\\sum\_\{t^\{\\prime\}\\in\\mathcal\{V\}\}q\_\{d\}\(t^\{\\prime\}\)\\rho\_\{d\}\(t^\{\\prime\}\\mid c\_\{d\}\)^\{\\lambda\_\{d\}\}\.
Thus, the optimal surrogate takes a product\-of\-experts form, combining the calibrated marginal signalqdq\_\{d\}with a prefix\-conditioned correctionρd\\rho\_\{d\}\.

## Appendix CPRESTO on Hybrid dLLMs at Quadratic Self Speculation Decoding

We provide a detailed walk\-through of how PRESTO is instantiated on Hybrid dLLMs at quadratic self speculation decoding mode, corresponding to Figure[3](https://arxiv.org/html/2607.22634#S4.F3)\. Unlike the standard speculative decoding setup where a small drafter proposes tokens for a larger verifier, self\-speculative dLLMs use*the same dLLM*as both drafter and verifier: a single forward pass simultaneously verifies previously drafted tokens and speculates new ones at trailing mask positions\. PRESTO turns this single\-model pipeline into a tree\-based drafter without any additional model\.

At the start of roundrr, the input sequence consists of three contiguous regions:

- •acommitted prefixof accepted tokens from previous rounds \(green tokens with ✓ in Figure[3](https://arxiv.org/html/2607.22634#S4.F3), e\.g\.C, D\);
- •averification regioncontaining the draft tokens proposed in roundr−1r\-1\(e\.g\.E, F, G\), some of which will be accepted and some rejected;
- •aspeculation regionof mask tokens \(M\) appended to the right, at which the dLLM will produce new draft candidates\.

A single forward pass over this composite input yields outputs that serve both roles below\.

#### \(I\) AR verification of the previous draft tree\.

On the verification region, PRESTO uses the exact KV cache of the committed prefix together with an*AR tree attention mask*that restricts each draft token to attend only to its ancestors in the previous round’s draft tree\. This lets a single forward verify all paths of the tree in parallel: each path is checked left\-to\-right against the model’s predictions, and the longest accepted prefix across paths is committed \(D→\\toFin Figure[3](https://arxiv.org/html/2607.22634#S4.F3), whereEon the sibling path andGon the same path are rejected\)\. The position immediately after the last accepted token is filled with the model’s own prediction \(F’,G’,H’in Figure[3](https://arxiv.org/html/2607.22634#S4.F3)\), and the accepted tokens become part of the committed prefix for roundr\+1r\+1\.

#### \(II\) Diffusion\-based speculation along the accepted path\.

The speculation region is structured as a set of mask slots attached to each draft token from the previous round, mirroring the previous round’s draft tree \(Figure[3](https://arxiv.org/html/2607.22634#S4.F3)shows two candidate paths under verification,D→\\toEandD→\\toF→\\toG\)\. The dLLM denoises all of these mask positions in parallel within the same forward, but only the mask slots hanging off the*accepted*path are used for the next round’s drafting\. In Figure[3](https://arxiv.org/html/2607.22634#S4.F3),Fis accepted andGis rejected, so PRESTO reads top\-bbcandidates from the mask slots followingF\(b=2b=2in the figure, e\.g\.\{H1,H2\}\\\{H\_\{1\},H\_\{2\}\\\}at the first such position,\{I1,I2\}\\\{I\_\{1\},I\_\{2\}\\\}at the next, and so on\)\. The mask slots attached to rejected branches are produced by the same forward but discarded \(marked “Unused” in Figure[3](https://arxiv.org/html/2607.22634#S4.F3)\), since their predictions were conditioned on a context that is no longer part of the committed prefix\.

#### \(III\) Draft tree construction\.

The per\-position top\-bbcandidates are assembled into a draft tree following the priority\-based expansion of Section[4\.2](https://arxiv.org/html/2607.22634#S4.SS2)\. The root corresponds to the first speculation position, and at depthddeach node has up tobbchildren drawn from the top\-bbcandidates at depthd\+1d\+1\. Each child inherits its parent’s cumulative path scoreSSand adds the local token scoresd\+1,t​\(c\)s\_\{d\+1,t\}\(c\), so scores are computed incrementally with no redundant work\. The frontier is expanded under beam search with global retention until the tree containsBBnodes, after which the top\-BBpaths underSSare kept \(Algorithm[1](https://arxiv.org/html/2607.22634#alg1)\)\. Crucially, the candidate logits at every depth come from the*same*forward pass that performed verification, so tree construction itself adds no additional dLLM forwards\.

#### \(IV\) Flattening for the next forward pass\.

The draft tree is then linearised into the input sequence for roundr\+1r\+1\(bottom of Figure[3](https://arxiv.org/html/2607.22634#S4.F3)\)\. Each draft token is followed by a block of mask tokens that will host the next round’s speculation, while the tree’s parent–child structure is encoded in the AR tree attention mask used during verification\. The result is a single input on which the next forward pass again performs verification and speculation jointly, completing the cycle\.

## Appendix DEnd\-to\-End Throughput Results for dFlash

Table 5:Average absolute decoding throughput \(tokens/s\) on Qwen3 models with thinking mode disabled\. Speedup is measured against autoregressive decoding on torch implementation with single batch size\.
## Appendix EEnd\-to\-End Throughput Results for Nemotron\-Labs\-Diffusion

Table 6:Absolute decoding throughput \(tokens/s\), average acceptance length \(τ\\tau\), and average throughput improvement on Nemotron\-Labs\-Diffusion\-8B under linear self\-speculative decoding with a maximum generation length of 2048 tokens based on SGLang implementation\.
## Appendix FEntropy\-adaptiveλd\\lambda\_\{d\}

In addition to a fixed interpolation coefficientλd\\lambda\_\{d\}, we explored an entropy\-adaptive variant that scales the prefix correction by the uncertainty of the diffusion drafter\. Let

Hd=−∑w∈𝒱qd​\(w\)​log⁡qd​\(w\)H\_\{d\}=\-\\sum\_\{w\\in\\mathcal\{V\}\}q\_\{d\}\(w\)\\log q\_\{d\}\(w\)denote the entropy of the diffusion marginal at depthdd\. We define

λd​\(Hd\)=β⋅HdHd\+C,\\lambda\_\{d\}\(H\_\{d\}\)=\\beta\\cdot\\frac\{H\_\{d\}\}\{H\_\{d\}\+C\},whereβ\>0\\beta\>0controls the maximum correction strength andC\>0C\>0is a stabilising constant, and substitute this into the token\-level score from Section 4:

sd,t​\(cd\)=log⁡qd​\(t\)\+λd​\(Hd\)​log⁡ρd​\(t∣cd\)\.s\_\{d,t\}\(c\_\{d\}\)=\\log q\_\{d\}\(t\)\+\\lambda\_\{d\}\(H\_\{d\}\)\\,\\log\\rho\_\{d\}\(t\\mid c\_\{d\}\)\.Under this schedule, low\-entropy \(high\-confidence\) positions assign small weight to the prefix correction and rely primarily on the calibrated marginalqdq\_\{d\}, while high\-entropy positions up\-weight the prefix\-conditioned signalρd\\rho\_\{d\}\. The fixed\-λd\\lambda\_\{d\}variant used in our main experiments is recovered by replacingλd​\(Hd\)\\lambda\_\{d\}\(H\_\{d\}\)with a constant\.