Unlocking Lossless Speedups in LLMs via Discrete Diffusion

Unlocking Lossless Speedups in LLMs via Discrete Diffusion Subham Sekhar Sahoo∗,†,1 , Lingjie Chen†,1,2 , Khiem Pham†,1,3 , Jonathan Geuter†,1,4 , Chaitanya Dwivedi1 , Varad Pimpalkhute1 , Yash Akhauri1 , Alexander Moreno1 , Mikhail Yurochkin1 , Zhenting Wang1 , Mostafa Elhoushi5 , Nolan Dey5 , Shane Bergsma5 , Joel Hestness5 , John Thickstun3 , Eric Xing1 , Zhengzhong Liu1 Institue of Foundation Models, 2 University of Illinois Urbana-Champaign, 3 Cornell Tech 4 Harvard University 5 Cerebras Systems † Core Contributors ∗ Correspondence to subham.sahoo@mbzuai.ac.ae

Large Language Models (LLMs) owe much of their success to next-token prediction (NTP), but their autoregressive (AR) structure requires slow, sequential token generation. To overcome this bottleneck, we introduce diffusion-augmented LLMs, a new class of models that defines an AR model distribution while using diffusion to draw multiple tokens in parallel from that distribution. We decouple the parameters of these models into two sets: AR weights, trained using the standard NTP objective, and lightweight diffusion weights, trained to generate multiple tokens simultaneously. The diffusion weights are learned through a simple Diffusion Distillation phase that adds negligible overhead to existing LLM training pipelines. We also introduce Ψ-Spec, a family of samplers that enables lossless acceleration and inference-time scaling at a fixed context length. Unlike speculative decoding, our method requires no separate draft model. Unlike diffusion LLMs (d-LLMs), it accelerates generation without sacrificing the quality of the underlying AR model. The resulting models, called Uno, can be trained from scratch or built by augmenting existing open-weight AR LLMs. Uno achieves higher throughput than leading speculative-decoding methods at every evaluated batch size and delivers up to 3× speedups over the base AR model, including at the largest batch size supported by the device. Notably, our 8B Uno model outperforms the leading open d-LLM, the 26B DiffusionGemma, and the proprietary Mercury 2 across all evaluated benchmarks in agentic tool use, coding, and long-context reasoning. We release code and checkpoints on: https://s-sahoo.com/uno

Diffusion-Augmented LLM Training DROP-IN PHASE

NTP Training

Diffusion Distillation

Pre-Training & Instruction-following

RL Post-training

Lossless Inference Acceleration Uno, AR (Ours)

Mercury 2 (Inception Labs)

90

71

68

68

68

40

40

Uno

AR

3000

58

58

AR

90

Improved Reasoning & Preference alignment Diffusion Gemma (Google)

Uno

Speed

5000

AR

Uno

AR

Uno

1000

AR

36

Uno

arXiv:2609.04010v1 [cs.LG] 3 Sep 2026

1

20

27 15

System Throughput: Tokens / Sec

18 N/A

Agentic Tool Use

Long Context Reasoning

Agentic Terminal Use

Agentic Coding

(Tau2 Banking)

(AA-LCR)

(Terminal-Bench v2.1)

(SWE-bench Verified)

Figure 1: (Top) Training overview for diffusion-augmented LLMs. Gray cells indicate AR-weight training, while the blue cell indicates diffusion-weight training. (Bottom Left) System throughput of Uno, the base AR model, and the baselines; see Sec. 5.1 for details. (Bottom Right) Performance across agentic and long-context reasoning benchmarks. 1

1

Introduction

Large language models (LLMs) are increasingly showing human-level capabilities across coding, mathematics, and complex reasoning (KimiTeam et al., 2026; GLM-5-Team et al., 2026). These capabilities arise from a simple yet highly effective training objective: next-token prediction (NTP) over massive text corpora (Brown et al., 2020). A key limitation of this objective is that inference is inefficient. By training models to predict only the next token, it produces autoregressive (AR) systems that generate just one token per decoding step. This sequential constraint becomes increasingly costly as reasoning traces grow longer, increasing serving latency and slowing reinforcementlearning (RL) post-training, where rollout generation dominates runtime (Kim et al., 2026). Sequential decoding is poorly suited to both natural language and modern accelerators. Language contains predictable collocations and formulaic sequences that could be generated together in blocks (Church & Hanks, 1990; Biber, 2009), but standard LLMs cannot exploit this redundancy. Moreover, decoding is often memory bound, especially at long context lengths, because moving model weights and key-value states limits inference speed and leaves GPUs underutilized (Verma & Vaidya, 2023). Predicting multiple tokens per step can amortize these memory transfers and better use parallel compute. Existing approaches only partially address this opportunity. Speculative decoding accelerates generation by verifying tokens proposed by a smaller draft model (Leviathan et al., 2023; Chen et al., 2023), but its gains depend on an efficient, well-aligned drafter and require maintaining a separate model. Discrete diffusion models (Austin et al., 2021a; Sahoo, 2026) support parallel generation natively and have succeeded in biological domains (Sahoo et al., 2024b; Schiff et al., 2025; Tang et al., 2025; Lee et al., 2025), yet leading d-LLMs (d-LLMs) still face a quality-speed tradeoff relative to AR models (Sahoo et al., 2025b, 2026; Bie et al., 2026; DiffusionGemmaTeam et al., 2026). Their speedups also vanish at large inference batch sizes (Wu et al., 2025; Fu et al., 2026; DiffusionGemmaTeam et al., 2026). The (lossy) speedups offered by d-LLMs at low batch sizes can therefore be insufficient in practice, as agentic workloads now predominate in contemporary LLM applications (Google Cloud, 2025; Amazon Web Services, 2026). In these workloads, a single request can trigger parallel agents, branching trajectories, tool calls, and retries, while serving systems batch calls across requests to improve accelerator utilization. Batch-size-one latency therefore captures a narrow operating regime and may overstate speedups that diminish under concurrency. Thus, practical acceleration should be evaluated at realistic batch sizes. Our core idea is simple: define a high-quality AR distribution, then learn to sample multiple tokens in parallel from that same distribution. We realize this idea through diffusion-augmented LLMs (Sec. 3), which use a single architecture with two decoupled sets of weights: one governing response quality and the other optimized for generation speed. Unlike multi-token prediction (MTP) approaches (Cai et al., 2024; Gloeckle et al., 2024), which modify the LLM architecture by adding prediction heads for future tokens, our method augments each layer in the LLM with lightweight diffusion weights alongside the standard AR weights. Our training pipeline first trains the AR weights through NTP loss, then freezes them and learns the diffusion weights through Diffusion Distillation (Sec. 3) to generate token blocks in parallel with negligible training overhead. Our Ψ-Spec sampler (Sec. 4) performs parallel prediction from the AR distribution. The resulting model is a drop-in alternative to speculative decoding and self-speculative decoding, requiring neither a separate draft model nor lossy AR-to-diffusion conversion. We call our models Uno because they unify AR and diffusion weights within one architecture. Across our evaluations, Uno improves the speed-quality frontier over the speculative-decoding methods EAGLE-3 (Li et al., 2026b) and DFlash (Chen et al., 2026), the open-weight d-LLMs Nemotron-Labs-Diffusion (Fu et al., 2026) and DiffusionGemma (DiffusionGemmaTeam et al., 2026), and the proprietary d-LLM Mercury 2 (Ermon, 2026). Its speedup over the base AR model persists at every evaluated batch size, including the largest batch that fits on the device with the AR weights, supporting both low-latency generation and high-throughput serving. Our contributions are: 1. We introduce diffusion-augmented LLMs, which decouple the parameters that determine response quality from lightweight parameters optimized for generation speed, together with a corresponding training pipeline (Sec. 3). 2

2. We propose Ψ-Spec samplers to sample from diffusion-augmented LLMs. These samplers enable both lossless, AR-verified acceleration and inference-time scaling through additional denoising at a fixed context length; see Sec. 4. 3. We show how to train diffusion-augmented LLMs from scratch (Sec. 5.1) or create them by augmenting an existing open-weight AR LLM (Sec. 5.2). 4. We show that Uno achieves up to a 2× speedup at the largest batch size supported by the base AR model, accelerating both inference and end-to-end RL training (Sec. 5.1).

2

Background

Notation We denote scalar discrete random variables with K categories as ‘one-hot’ column vectors and define V = {v ∈ {0, 1}K ∶ ∑K i=1 vi = 1} as the set of all such vectors. Define Cat(⋅ ; π) as the categorical distribution over K classes with probabilities given by π ∈ ∆, where ∆ denotes the K − 1-simplex. We also assume that the K-th category corresponds to a special [MASK] token and let m ∈ V be the one-hot vector for this mask, i.e., mK = 1. Additionally, let 1 = {1}K , and ⟨a, b⟩ and a ⊙ b respectively denote the dot and Hadamard products between two vectors a and b. We use x ∈ V L to represent clean data which is a length-L sequence containing no mask tokens, and let xℓ denote its ℓth element where ℓ refers to the token index. Under our notation, each xℓ is a onehot vector. The notation xm∶n denotes the subsequence from index m to index n, inclusive, while x<ℓ ∶= x1∶ℓ−1 denotes the prefix consisting of the first ℓ − 1 tokens. Finally, [⋅, ⋅] ∶ V m × V n → V m+n denotes the operator that concatenates sequences of lengths m and n. Colors are used throughout the paper solely for emphasis and carry no mathematical meaning.

2.1

Autoregressive Models

Consider a length-L sequence x ∈ V L ∼ qdata drawn from the data distribution qdata . Autoregressive (AR) language models factorize the joint distribution using the chain rule: L

log pθ (x) = ∑ log pθ (xℓ ∣ x<ℓ ). ℓ=1 L

L

Here, pθ ∶ V → ∆ is typically implemented with a causal Transformer (Vaswani et al., 2017) with parameters θ. This factorization leads to strong likelihood modeling and enables efficient inference primitives such as KV caching. However, AR generation is intrinsically sequential: ℓth token can only be generated after the prefix x<ℓ has been generated. As a result, producing an output of length L requires L causal decoding decisions. Speculative Decoding Speculative decoding (Leviathan et al., 2023; Chen et al., 2023) uses a smaller draft model to autoregressively generate a block of candidate tokens, while a larger base model verifies them. The key observation is that although sampling from the base model is inherently sequential and expensive, evaluating the likelihood of an entire candidate sequence can be performed efficiently in parallel. Speculative decoding exploits this asymmetry by generating candidate tokens with the draft model that can be sampled more efficiently, and verifying them via rejection sampling in a single forward pass, accepting the longest valid prefix. This accelerates inference while exactly preserving the target distribution.

2.2

Discrete Diffusion Language Models

Discrete diffusion corrupts a clean sequence x ∈ V L ∼ qdata into a simple prior π L and learns to reverse the corruption. We use the interpolating forward process from Sahoo et al. (2024a): zℓt ∼ qtℓ (⋅ ∣ x; π) ∶= Cat(⋅ ; αt xℓ + (1 − αt )π) , (1) where zt denotes the noisy sequence at time step t ∈ [0, 1], and the noise schedule αt decreases monotonically from α0 = 1 at clean data to α1 = 0 at the prior. A denoising model xθ ∶ V L ×[0, 1] → ∆L , parameterized by θ, predicts the clean tokens from zt and t. Masked diffusion uses π = m, while uniform-state diffusion uses π = 1/K. We use the latter because it natively supports selfcorrection, few-step generation (Sahoo et al., 2025a), and better inference-time scaling than masked diffusion models (Deschenaux et al., 2026). 3

Sampling For 0 ≤ s < t, let qs∣t denote the reverse posterior that maps zt to a less noisy state zs . The Ψ-samplers proposed by Deschenaux et al. (2026) form a family of discrete-diffusion samplers that generalize posterior (ancestral) sampling by incorporating predictor-corrector capabilities, thereby producing higher-quality samples. Their practical token-wise transition kernel is given by ℓ ℓ [Ψθs∣t (⋅ ∣ zt ; xθ (zt ), π)]ℓ = κt qs∣t (⋅ ∣ zt , xθ (zt ); π) + (1 − κt ) [αs q0∣t (⋅ ∣ zt , xθ (zt ); π) + (1 − αs )π] , (2)

where κt ∈ [0, 1] controls the strength of the correction. The exact functional form of (2) for MDMs and USDMs, see Suppl. A.4. Discrete Consistency Distillation Discrete Consistency Distillation (DCD) compresses a multistep uniform-state diffusion process into a few-step generator (Sahoo et al., 2025a). It constructs a deterministic discrete trajectory from z1 ∼ 1L /K to x using the Gaussian diffusion process underlying the uniform-state discrete diffusion process,. For adjacent states zt and zs , with s < t, the student distribution xℓθ (zt , t) is trained to match the teacher distribution xℓθ0 (zs , s); ∀ℓ ∈ [L]: L

LDCD (θ; θ0 ) = ∑ DKL (xℓθ (zt , t) ∥ xℓθ0 (zs , s)) . ℓ=1

Increasing the gap between s and t across distillation rounds teaches the student to take larger denoising steps. Refer Suppl. A.5 for further details.

3

Diffusion-augmented LLMs

We propose a new framework for designing LLMs that decouples generation quality from generation speed by augmenting each layer of a standard AR model with a separate set of diffusion weights dedicated to parallel token generation. We call the resulting model a diffusion-augmented LLM. The key distinction is that each layer contains two sets of weights: AR weights, trained to generate tokens autoregressively and determine response quality, and diffusion weights, trained to generate tokens in parallel. This separation preserves the established AR training pipeline while introducing a dedicated phase for training the diffusion weights to accelerate inference. At generation time, both sets of weights draft tokens in parallel, after which the AR weights alone verify the drafts. Our sampler provably preserves the AR model’s output distribution (Sec. 4), enabling lossless acceleration without sacrificing response quality. The framework can be applied to any causal autoregressive neural network, including causal Transformers (Vaswani et al., 2017) and State-Space Models (SSMs; Gu et al., 2020; Gu & Dao, 2023). Because of this separation, the framework can also augment existing open-weight LLMs with parallel generation capabilities, providing lossless speedups without retraining their base parameters (Sec. 5.2). The remainder of this section develops the framework in three stages. We first describe the roles of the AR and diffusion weights in a Diffusion-augmented LLM (Sec. 3.1). We then present the Diffusion Distillation Phase (Sec. 3.2), which trains the diffusion weights using a one-step blockdenoising formulation and associated training objectives. Finally, we discuss the training curriculum for two settings: accelerating inference alone and accelerating both RL rollouts and inference (Sec. 3.3).

3.1

Autoregressive and Diffusion Pathways

Each layer of a diffusion-augmented LLM contains a set of autoregressive (AR) weights and a set of diffusion weights. These two sets of weights follow different training procedures, as described below. Autoregressive Weights We denote the base AR weights by θAR . These parameters are trained using the standard LLM recipe of Pre-training (Brown et al., 2020; Hoffmann et al., 2022; Grattafiori et al., 2024), supervised fine-tuning (SFT; Wei et al., 2021; Chung et al., 2024), and RL posttraining (Ouyang et al., 2022; Shao et al., 2024; Guo et al., 2025). Given a length-L sequence x ∈ V L , the model xθAR ∶ V L → ∆L is causally masked, so its output at position ℓ observes only the prefix x<ℓ , the distribution of the ℓth token is modeled as: pθAR (xℓ ∣ x<ℓ ) = xℓ−1 1 < ℓ ≤ L. θAR (x), We assume that the first token in x is a special <BOS> token that requires no modeling. 4

Diffusion Weights The diffusion weights θ∆ are used solely to accelerate generation. As described in Sec. 4, our sampler uses the diffusion pathway to draft tokens and performs rejection sampling against the AR distribution defined by θAR . We parameterize the diffusion weights as LoRA (Low RAnk) adapters (Hu et al., 2022): each AR weight matrix has an associated LoRA weight. Thus, the diffusion pathway uses θAR + θ∆ to generate draft tokens, whereas the verification pathway uses θAR , with the two distributions coupled through their shared base parameters. We train only θ∆ during the Diffusion Distillation Phase while keeping θAR frozen; see Sec. 3.2. For effective rejection sampling, the draft distribution must remain closely coupled to the verification distribution. Parameterizing the draft pathway as a LoRA adaptation of the base model promotes this coupling while adding minimal inference-time memory overhead compared with a separate set of full-sized diffusion weights. This design retains the lossless verification of standard speculative decoding without requiring a separate draft model. At the same time, like self-speculative decoding, it tightly couples the drafting and verification pathways, but does so while remaining lossless.. Diffusion training teaches the LLM to refine an entire sequence in parallel. Specifically, given a noisy sequence (zt ∈ V L )t∈[0,1] obtained by corrupting a clean sequence x ∈ V L according to (1), the model is trained to recover x. Throughout this work, the LLM retains the NTP parameterization of standard autoregressive LLMs: the logits at each position predict the subsequent clean token. This parameterization differs from that of conventional diffusion language models (Austin et al., 2021a; Lou et al., 2024; Sahoo et al., 2024b), which typically predict the clean token at the same position. Accordingly, the distribution of the clean token xℓ is: pθAR ,θ∆ (xℓ ∣ zt ) = xℓ−1 θAR ,θ∆ (zt ),

3.2

1 < ℓ ≤ L.

Diffusion Distillation Phase

We train the diffusion parameters θ∆ to draft a block of tokens in parallel that the autoregressive (AR) model with parameters θAR would generate sequentially. Specifically, our goal is to match the single-step diffusion distribution pθAR ,θ∆ (⋅) over a token block to the AR distribution pθAR (⋅) over the same block. Discrete Consistency Distillation (DCD; Sec. 2.2) is well suited to this task because it can distill a multistep diffusion process into a few-step generator while closely approximating the distribution induced by the original process. In what follows, we first describe how we adapt DCD to approximate the AR distribution through single-step diffusion. We then extend the DCD training objective, namely the single-step distillation loss LDCD , with another auxiliary loss term: LTV to encourage longer diffusion drafts to be accepted by the AR verifier. The resulting objective is: L(θ∆ ; θAR , α, β) = Ex∼D,z1 ∼πL [αLDCD (θ∆ ; θAR , x, z1 ) + βLTV (θ∆ ; θAR , x, z1 )] ,

(3)

where z1 is the fully corrupt sequence. Although α = 0, β = 1 yields the largest speedups, training first with α = β = 1 and then switching to α = 0, β = 1 accelerates convergence. We ablate both loss terms in Sec. 5.2.4. We describe each term below. One-step Distillation Standard DCD constructs a multistep denoising trajectory by simulating intermediate states of the underlying PF-ODE. At LLM scale, repeatedly constructing and storing these intermediate Gaussian latents is prohibitively expensive. We therefore distill the entire trajectory into a single denoising step that maps a fully corrupted input z1 ∼ π L directly to the clean sequence x. The frozen base model defines the autoregressive teacher distribution xθAR , and the adapted model defines the student distribution xθAR ,θ∆ . We train the student to match the teacher directly, eliminating the need to simulate or materialize any intermediate PF-ODE states. Block Diffusion Recovering an entire long sequence from a fully corrupted input in a single denoising step is prohibitively difficult; we therefore perform one-step denoising blockwise. We partition the clean sequence x and the fully corrupt sequence z1 into N blocks of size B. We use x(b) (b) and z1 to denote the bth blocks of x and z1 , respectively. For each block, we train θ∆ to match (b) the student diffusion distribution over the clean block x(b) , conditioned on the noisy block z1 and (<b) the preceding clean context x , to the teacher AR distribution over the same block, conditioned on the preceding clean context. We compute the teacher and student predictions in a single forward pass over the concatenated sequence [x, z1 ]. Specifically, we use a block-causal attention mask 5

(b)

(b)

that permits causal attention within x and within each noisy block z1 . Tokens in z1 additionally attend to all preceding clean blocks x(<b) . The main challenge is to compute the teacher logits at positions in x using only θAR , while computing the student logits at positions in z1 using both θAR and θ∆ . We achieve this using gated LoRA (Samragh et al., 2025), which disables the adapters at clean-sequence positions and enables them at noisy-sequence positions. Consequently, θAR produces the teacher logits, whereas θAR and θ∆ jointly produce the student logits. The resulting blockwise DCD objective is: N

B

(N +b,ℓ)

(b,ℓ)

LDCD (θ∆ ; θAR , x, z1 ) = ∑ ∑ DKL (xθ∆ ,θAR ([x, z1 ])∥xθAR ([x, z1 ])) . b=1 ℓ=1

The base parameters θAR remain frozen throughout fine-tuning, and only the LoRA parameters θ∆ are updated. Total Variation Loss Our sampler (Sec. 4) performs rejection sampling against the AR distribution and retains only the longest consecutive prefix of the diffusion prediction. Sampling efficiency therefore depends on the length of this retained prefix. To increase its expected length, we minimize the blockwise Total Variation (TV) distance between the diffusion and the AR distributions following Corollary 3.6 of Leviathan et al. (2023): N

B

(N +b,ℓ)

(b,ℓ)

LTV (θ∆ ; θAR , x, z1 ) = ∑ ∑ ∣xθ∆ ,θAR ([x, z1 ]) − xθAR ([x, z1 ])∣ . b=1 ℓ=1

Minimizing this objective increases the probability of consecutive tokens being accepted, and, consequently, increases the expected length of the accepted draft prefix.

3.3

Practical Considerations

The AR weights θAR and diffusion adapters θ∆ serve complementary roles and are optimized using separate objectives, making their training order an important practical consideration. In particular, Diffusion Distillation can be performed either after all AR training is complete or before RL posttraining so that the resulting adapters can accelerate rollout generation. We therefore consider two training regimes, depending on whether diffusion-based generation is used only during inference or during both RL post-training and inference. Faster Inference Only If the sole objective is faster inference, we first complete autoregressive pre-training and post-training, or simply initialize them with an Open-weights LLM. We then freeze the AR weights θAR and train θ∆ using Diffusion Distillation (Sec. 3.2). Faster RL Training & Faster Inference We demonstrate in Sec. 5 that our method provides significant generation speedups wrt the base AR model across all batch sizes, and can therefore accelerate RL post-training due to faster rollouts. In this setting, we first apply diffusion distillation after supervised fine-tuning and before RL, as illustrated in Figure 1 (top). The resulting adapters can then accelerate rollouts during RL. Standard RL policy-optimization recipes, including PPO and GRPO-based methods (Ouyang et al., 2022; Shao et al., 2024; Guo et al., 2025; Yu et al., 2026a), update θAR but not the θ∆ . One might therefore expect that, as θAR changes, the draft distribution will drift away from the verifier distribution, reducing the acceptance rate and eroding the speedup. Surprisingly, we show in Sec. 5.1 that the speedup is retained even as the AR weights are trained using RL.

4

Ψ-Speculative Sampler

Our goal is to sample from the AR distribution defined by θAR while drawing multiple tokens in parallel. We therefore introduce the Ψ-Speculative sampler (Ψ-Spec), which uses the diffusion pathway to propose a block of tokens in parallel (Sec. 4.1) and performs rejection sampling against the base AR distribution to accept the longest valid prefix (Sec. 4.2). 6

4.1

Diffusion Sampler

Let x ∈ V L denote a partially generated sequence of length L. To generate a block of B tokens, we first append B − 1 random tokens sampled from the prior, z1 ∼ π B−1 . Denoising then proceeds from t = 1 to t = 0, with zt denoting the partially denoised sequence at time t. At each step, the denoiser xθ (⋅ ; x<L ) ∶ V B → ∆B operates on the block [xL , zt ] ∈ V B , while the activations for the preceding sequence x<L remain in the KV cache. In Sec. 4.1.1, we define the proposal distribution over a block of tokens, and in Sec. 4.1.2, we describe how candidates are sampled from this distribution. 4.1.1

Diffusion Proposal Distribution

Given a noisy block zt , we use the Ψ-Spec transition in (2) to sample a less noisy block zs , where s < t. Because the denoiser uses the NTP parameterization, the clean-token distributions used to L construct zs are given by x1∶B−1 θAR ,θ∆ ([x , zt ]). The input contains both clean and corrupted tokens, while the diffusion weights are trained only on corrupted tokens. To avoid distribution shift, we compute the logits for the first, clean position using only the base AR weights, and those for the noisy positions using both the AR and diffusion weights. We achieve this in a single forward pass using the gated LoRA technique of Samragh et al. (2025). Let Ψ0 denote the distribution induced by the Ψ-sampler at t = 0. At the final step, we sample the first token using only the base AR weights: 1∶B−1 (4) z10 ∼ Ψℓ=1 0 (⋅ ∣ zt ; xθAR (⋅), π) . We sample the remaining B − 2 tokens from the following joint distribution, discussed in Sec. 4.1.2: B

z2∶B ∼ ∏ Ψℓ0 (⋅ ∣ zt ; x1∶B−1 θAR ,θ∆ (⋅), π) . 0

(5)

ℓ=2

4.1.2

Sampling Candidates

ℓ Given the proposal distribution in (5), we construct a candidate set C = {c ∶ c ∼ ∏B ℓ=2 Ψ0 }. The number of draft candidates determine the tradeoff between acceptance length and verification cost. Sampling more candidates increases the probability of accepting a longer prefix, but also increases the cost of verification. We therefore select the largest candidate set that can be verified without reducing inference throughput, subject to the available serving batch size.

Linear Sampler (System Throughput Optimized) The simplest approach samples each draft token directly from its marginal distribution in (5), producing a single candidate sequence. At high batch sizes, inference can become compute-bound, leaving little spare compute for verifying additional candidates. The linear sampler is therefore well suited to maximizing aggregate system throughput in this regime. Tree Sampler (Single-user Throughput Optimized) At low batch sizes, inference is typically memory-bound, leaving substantial compute capacity underutilized. We exploit this spare compute by sampling multiple candidates and verifying them in parallel. Specifically, we use the tree-based sampling procedure of Cai et al. (2024), which selects the top K tokens at each position within a block. Rather than evaluating all K B−1 candidate sequences, we adopt the candidate-pruning strategy of Li et al. (2026b) which ranks candidates by their log-probabilities and retains the top V ∈ Z+ prefixes. Thus, (B, K, V ) are the hyperparameters of the tree sampler, corresponding to the block size, branching factor, and prefix budget, respectively.

4.2

One-step Diffusion w/ AR Verification

The central goal of this work is to accelerate generation. We therefore draft an entire block of B tokens using a single diffusion forward pass. In Suppl. B.2, we show that, for single-step diffusion generation, Eqns. (4, 5) reduce to z10 ∼ x1θAR (⋅), B

z2∶B ∼ pdraft = ∏ xℓθAR ,θ∆ (⋅). 0 ℓ=2

7

From pdraft , we sample candidates C = {c ∶ c ∼ pdraft } and apply the standard speculative-decoding rejection correction (Leviathan et al., 2023), retaining the longest prefix accepted by the base AR model. This preserves the base model’s target distribution. For ∣C∣ > 1, candidates are verified concurrently as a prefix tree using tree attention (Cai et al., 2024). Because θAR remains frozen and only the diffusion LoRA adapters are trained, the base AR model provides an unchanged verifier and enables lossless speculative speedups. The complete sampling algorithm is provided in Algo. 1. Tokens-Per-Forward-pass (TPF) During drafting, the first token is generated using the base AR weights and therefore matches the verifier distribution exactly, so it is always accepted. If a subsequent draft token is rejected, the verifier samples a replacement from the renormalized residual distribution, as in standard speculative decoding. Consequently, even an immediate rejection produces two output tokens: the always-accepted first token and the verifier-sampled replacement token. If all B draft tokens are accepted, the verifier additionally samples one token from the logits following the final draft token, producing B + 1 output tokens. Because each iteration requires two forward passes, one for drafting and one for verification, the TPF is bounded by 1 ≤ TPF ≤ B+1 . 2 We provide support for both the Linear and Tree samplers in Nano-vLLM (Yu, 2025) and SGLang (Zheng et al., 2024). All experiments reported in this paper were conducted using our Nano-vLLM implementation.

4.3

Inference-Time Scaling

Ψ-Spec introduces an additional axis for inference-time scaling in LLMs by allowing the number of denoising steps T to exceed the number of drafted tokens B. Deschenaux et al. (2026) show that, for 0 ≤ κt < 1, increasing T consistently improves sample quality. The central question is whether this improvement can eventually surpass the quality of AR generation. If so, AR verification should be disabled, since it would constrain the final output to the lower quality level of the AR model and prevent the gains from additional denoising from being realized. Unlike conventional inferencetime scaling methods, this approach allocates additional computation without increasing the context length. We leave a systematic exploration of this quality-compute tradeoff to future work.

5

Experiments

We train diffusion-augmented LLMs in two settings. In the first, we train the AR weights from scratch on proprietary data and train the diffusion weights on data drawn from the same distribution. In the second, we augment an open-weight AR model, Qwen3-8B (Yang et al., 2025), without access to its original training data. We instead train the diffusion weights on a different data distribution using the open-source OpenThoughts dataset (Guha et al., 2025). Together, these experiments demonstrate the versatility of our method and show that the diffusion and AR weights need not be trained on the same data distribution. Throughput Analysis Throughput depends strongly on context length. However, existing evaluations (Wu et al., 2025; Fu et al., 2026) measure throughput on downstream tasks whose reasoning traces vary in length. This can misleadingly favor models that generate shorter traces. Following standard LLM evaluation practice1 , we instead evaluate every method using m random input tokens and a fixed output length of n tokens. For diffusion and speculative decoding methods, which draft a block of B tokens per forward pass, we first measure the average number of accepted tokens per n ⌉ decoding steps, constraining each forward pass (TPF) across all benchmarks. We then run ⌈ TPF method to accept TPF tokens per step on average. This ensures that all methods are evaluated at the same input and effective output lengths. Throughout the paper, we use m = 1024 and n = 8192 and refer to this setting as the “1K/8K throughput test”. Per-request Throughput vs. System Throughput We report throughput at batch size 1 and at the largest batch size that fits on a single H200 GPU. Batch-size-1 throughput is less representative of practical serving because agentic workloads often generate concurrent requests even while serving a single user as described in Sec. 1. We therefore also report system throughput, which is 1 https://nvidia.github.io/TensorRT-LLM/developer-guide/perf-benchmarking.html

8

the aggregate token generation rate at the maximum feasible batch size, which better reflects serving capacity and cost efficiency. For each regime, we select the TPF setting that maximizes throughput and use the 1K/8K test described above.

5.1

End to End Diffusion Augmented LLM Training

5.1.1

Setup

Benchmarks We evaluate on both agentic and non-agentic benchmarks. The agentic suite comprises the Telecom, Airline, and Retail domains of τ 3 -Bench (Shi et al., 2026), τ 2 -Bench (Barres et al., 2025) , Terminal-Bench v2.1 (The Terminal-Bench Team, 2026), and SWE-bench Verified (Jimenez et al., 2024). The non-agentic suite includes AA-Omniscience (Jackson et al., 2025), AA-LCR (Artificial Analysis, 2025), Humanity’s Last Exam (HLE; Phan et al., 2025) , GPQA-Diamond (Rein et al., 2024), GSM8K (Cobbe et al., 2021), MATH500 (Lightman et al., 2024), AIME 2024 and 2025 (Balunović et al., 2025), AIME 2026 (Dekoninck et al., 2026), and MBPP (Austin et al., 2021b). We use a context window of 262, 144 with 131, 072 maximum generation tokens, allowing for long prompts on agentic benchmarks, and sampling parameters temp = 1, top-p = 0.95, and top-k = 50. Metrics For each method, we report average pass@1 over multiple generations and average TPF; see Table 4 for details. We compute throughput using the “1k/8k throughput test,” averaging TPF across benchmarks. System throughput is measured using the largest batch size that fits on a single H200 GPU, whereas per-request throughput is measured at batch size 1. Baselines We compare against Nemotron-Labs-Diffusion-14B (Fu et al., 2026), DiffusionGemma-26B-A4B (DiffusionGemmaTeam et al., 2026), and Mercury 2 (Ermon, 2026). Nemotron-Labs-Diffusion is pretrained autoregressively on 1T tokens and then fine-tuned as a diffusion model on 300B tokens. It also uses masked diffusion with bidirectional attention within draft blocks. We evaluate the largest Nemotron-Labs-Diffusion variant because it performs better on benchmarks than the smaller variants. DiffusionGemma is a sparse 26B mixture-of-experts model with 4B active parameters, fine-tuned from Gemma 4 (DiffusionGemmaTeam et al., 2026). Since both Nemotron-Labs-Diffusion and DiffusionGemma modify the base AR parameters, they are lossy unlike Uno. We use the default baseline configurations: Nemotron-Labs-Diffusion uses Linear Self-Speculation with a block size of 32, a maximum generation length of 8192, a thinking-token budget of 6000, and temperatures of 0 for drafting and verification. DiffusionGemma uses a block size of 256, temperature 1, a maximum generation length of 131072, a context length of 162144, and an entropy-bounded diffusion sampler with a bound of 0.1 and thinking enabled. We also compare against Inception Labs’ closed-source d-LLM, Mercury 2, whose parameter count is undisclosed. 5.1.2

Uno: Diffusion Augmented LLM (Ours)

We first define the architecture of our model. The weights corresponding to every layer are the the AR weights θAR . Corresponding to every weight matrix, we introduce rank-128 LoRA adapters with LoRA-α = 256 which serve as the diffusion weights θ∆ . We share more details subsequently. LLM Architecture We use a dense decoder-only causal Transformer with 36 layers, hidden width 4, 096, SwiGLU-style MLP width 12, 288, 32 query heads, 8 KV heads, head dimension 128, grouped RMSNorm, RoPE (θ = 107 ), a 250, 624-token vocabulary, and a maximum configured context length of 524, 288 tokens. The model contains 6.95B transformer-body parameters plus approximately 2.05B parameters from the large untied input/output vocabulary matrices. Training AR weights The AR weights are trained on approximately 23T tokens on internal quality data with a staged context-length extension schedule. Pretraining uses 21.9T tokens at an 8K-token context length, with a peak learning rate of 3 × 10−4 under a warmup-stable-decay (WSD) configuration and a 3,000-step warmup. We then extend the context length to 32K, training on an additional 1.1T tokens while linearly decaying the learning rate from 3 × 10−4 to 3 × 10−5 , following a 1,250step warmup. Next, the context length is increased to 128K and the model is trained on 0.5T tokens 9

at a constant learning rate of 3 × 10−5 , with a 500-step warmup. Finally, the context length is extended to 512K and trained for 0.3T tokens with a constant learning rate of 3 × 10−5 and a 200-step warmup. Training Diffusion Weights Diffusion weights are implemented as rank-128 LoRA adapters with LoRA-α = 256. We train them on 7B tokens sampled randomly from the SFT training data while keeping the AR weights frozen. We employ curricula for context length and diffusion block size: 1.8B tokens at context length 16, 384, divided into increasing block sizes of two, four, and eight with 600M tokens each; followed by 5.2B tokens at context length 65, 536 and block size 8. We train on a global batch size of 128 with a WSD learning rate scheduler with 200 warmup steps and a peak learning rate of 5 × 10−5 . In (3), we set α = 0.01, β = 1. Training takes approximately 60 hours on 8 H200 nodes with 8 GPUs each. To sample from Uno, we use the Linear Sampler with B = 4 to maximize system throughput and the Tree Sampler with (B, K, V ) = (16, 32, 32) to maximize single-request throughput at batch size 1. 5.1.3

Results

Ψ-Spec Sampler Configurations We examine how TPF and throughput vary across Linear and Tree sampler configurations in Table 6. For the Linear sampler, TPF increases from 1.8 at B = 4 to 2.2 at B = 8, but reaches only 2.3 at B = 16; B = 4 achieves the highest system throughput. For the Tree sampler, with B = 16 and K = 32, V = 32 provides higher per-request throughput than V = 64, likely because the larger verification cost reduces throughput. Uno vs. Base AR Uno qualitatively matches the base AR model while achieving higher throughput at every batch size. In the 1k/8k throughput test, the largest batch size supported by the base AR model is 64, at which Uno is 1.5× faster (see Table 7). This speedup benefits multi-user serving and RL post-training. At batch size 1, Uno is approximately 2.2× faster. Takeaway 1: Uno matches the qualitative performance of Base AR while delivering over 2× higher throughput across all batch sizes. Comparison with Open Weights Diffusion Language Models As shown in Table 1, Uno outperforms the open-weight DiffusionGemma and Nemotron-Labs-Diffusion models on all tasks, with larger gains on agentic tasks. Using the TPF values in Table 1, we measure maximum system throughput at the largest feasible batch size and per-request throughput at batch size 1. Despite using full attention in every layer, Uno achieves higher system throughput (with Tree Sampler; (B, K, V ) = (16, 32, 32)) than both baselines. Notably, DiffusionGemma uses strided attention, with five local sliding-window self-attention layers for every global self-attention layer, yet remains slower at large batch sizes. DiffusionGemma is faster at batch size 1 but has substantially lower accuracy, while Nemotron-Labs-Diffusion trails Uno in both throughput and output quality. Although smaller blocks may improve throughput, further tuning offers no practical benefit given the baselines’ lower quality. Also, note that both baselines are slower than their respective base AR models at large batch sizes (Fu et al., 2026; DiffusionGemmaTeam et al., 2026). Takeaway 2: Uno outperforms open-weight d-LLMs on every benchmark, while also achieving the highest system throughput. Comparison with Proprietary Diffusion Language Models As shown in Table 1, the 8Bparameter Uno outperforms Mercury 2 on all Agentic Tool Use, Agentic Coding, and Long-Context Reasoning benchmarks, trailing only on one Science and Knowledge benchmark, potentially due to differences in model size and training data. More importantly, Mercury 22 reports a system throughput of 1,154 tokens/s for a 1K-token input at batch size 10. Uno achieves a maximum system throughput ∼ 4.6× higher, despite Mercury 2 running on substantially faster Blackwell GPUs. Moreover, Mercury 2 does not disclose its quantization, so it may also benefit from lower precision, whereas Uno and the other baselines use bfloat16. This result highlights the strength of diffusionaugmented LLMs. 2 https://artificialanalysis.ai/models/mercury-2?speed=output-speed-by-prompt-type#speed-tabs

10

Table 1: Accuracy and TPF of Uno, Nemotron-Labs-Diffusion, Mercury 2, and DiffusionGemma on agentic and non-agentic benchmarks. Mercury 2 results are from Artificial Analysis. Results for Nemotron-Labs-Diffusion and DiffusionGemma are computed from their open-source checkpoints (Table 5). For Uno, subscripts report “TPF1 / TPF2 ,” where TPF1 uses the system-throughputoptimal Linear sampler with B = 4, and TPF2 uses (B, K, V ) = (16, 32, 32), optimized for perrequest throughput. ∗ Reported by Artificial Analysis’s live tracker on August 30, 2026. Model size

Uno (8B)

Agentic Tool Use τ 3 Banking 25.81.8/2.7 τ 2 Telecom 90.11.7/2.1 τ 2 Retail 67.11.8/2.4 Terminal-Bench v2.1 39.62.1/2.7 Agentic Coding SWE-bench Verified 68.42.2/3.1 Long-Context Reasoning AA-LCR 68.01.8/2.6 Science and Knowledge Humanity’s Last Exam 18.61.8/2.8 77.12.0/4.2 GPQA-Diamond AA-Omniscience 14.31.7/3.1 Math GSM8K 95.41.9/2.7 MATH500 98.91.9/2.4 93.01.9/2.5 AIME-24 AIME-25 90.71.8/2.4 AIME-26 86.31.8/2.4 Coding MBPP 84.11.8/2.4 HumanEval 95.21.9/3.4 Avg. TPF System Throughput Per-request Throughput

1.9 / 2.7 5255 405

Mercury 2 Diffusion-Gemma Nemotron-Labs-Diffusion N/A (26B-A4B) (14B) 9N/A 71N/A – 27N/A

– 68.118.8 65.523.7 14.714.1

– 14.34.8 5.63.1 4.57.5

–

18.75.6

0.81.5

36N/A

19.710.8

7.31.1

16N/A 77N/A 20N/A

9.214.1 70.711.9 17.79.9

2.67.2 40.47.6 11.011.1

– – – – –

95.128.9 92.424.1 73.716.9 74.318.9 70.717.8

93.16.1 89.25.6 56.74.9 40.04.5 46.74.8

– –

80.115.5 95.128.2

73.85.3 84.87.5

N/A 1197∗ 769∗

17.56 1136 836

5.41 2794 290

Takeaway 3: Uno beats proprietary d-LLM Mercury 2, on all agentic, coding, and long-context benchmarks and achieves ∼ 4.6× higher maximum system throughput, despite running on slower hardware. Faster RL Training From the supervised fine-tuning (SFT) checkpoint, we trained four experts in mathematics, code generation, tool use, and web search using the DAPO reinforcement learning algorithm. During this stage, we updated only the base AR weights; the diffusion weights trained on the SFT checkpoint remained frozen and were used only to speed up RL rollouts. This yielded up to a 40% end-to-end training speedup, primarily for the mathematics and code experts. Gains were smaller for the tool-use and search experts because tool calls dominated their runtime. We will provide detailed results in the next revision. We then consolidated the four experts into a single model using ISO-Merger (RAM; Yuan et al., 2026), a data-free method that combines specialists trained from a shared base checkpoint without additional rollouts. As shown in Table 8, the diffusion adapters trained on the SFT checkpoint retain their speedups after RL post-training, with only a nominal 6% decrease in TPFs.

5.2

Open Weights LLMs

We show that Diffusion-augmented LLMs can initialize their autoregressive (AR) weights from the open-weight Qwen3-8B model while retaining the speedups enabled by diffusion weights trained on a different data distribution, specifically, the open-source OpenThoughts dataset. Although training Qwen3-8B on this dataset degrades its quality (see Table 9), the diffusion weights trained on the same nevertheless enable lossless speedups on the same benchmarks. 11

5.2.1

Setup

Benchmarks We evaluate on mathematical reasoning (GSM8K, MATH500, AIME-24, AIME25, and AIME-26), code generation: HumanEval (Chen et al., 2021), MBPP (Austin et al., 2021b), and LiveCodeBench v6 (Jain et al., 2025)); Science reasoning: GPQA and GPQA-Diamond (Rein et al., 2024); instruction following: IFEval with prompt-level strict accuracy (Zhou et al., 2023), and General Knowledge: MMLU-Pro (Wang et al., 2024). Lossless Baselines (Speculative Decoding Methods) Our primary lossless baselines are EAGLE3 (Li et al., 2026b), which uses an AR drafter, and DFlash (Chen et al., 2026), which uses a diffusion drafter. Both draft models are designed for Qwen3, and we evaluate their open-source checkpoints. EAGLE-3 adds a lightweight 0.40B-parameter drafter that sequentially predicts tokens from fused intermediate features of the target model. DFlash instead adds a 1.05B-parameter diffusion drafter that predicts multiple tokens in parallel, conditioned on target-model features. Thus, DFlash introduces roughly three times as many parameters as our method. Its training is also substantially more expensive. For block size B and sequence length L, DFlash requires a training context length of B ⋅ L, whereas our method always uses 2 ⋅ L, independent of B. This makes DFlash significantly more expensive to train than our method. We evaluate the open-source checkpoints for EAGLE-3 and DFlash; see Table 5. Refer Suppl. C.1.2 for sampler configurations. Lossy Baselines (Diffusion Methods) We also compare against lossy parallel-generation methods. Jacobi Forcing (Hu et al., 2025), SDAR (Cheng et al., 2026), OPDLM (Su et al., 2026), and I-DLM (Yu et al., 2026b) are fine-tuned from Qwen3; Fast-dLLM v2 (Wu et al., 2025) is fine-tuned from Qwen2.5; FLARE (Zhu et al., 2026) is fine-tuned from Qwen3.5; and LLaDA2.1-Flash (Bie et al., 2026) is trained from scratch. We evaluate the open-source checkpoints for Fast-dLLM v2 and SDAR; see Table 5. OPDLM, I-DLM, and FLARE are concurrent works. Metrics We report average pass@1 accuracy over multiple generations for each benchmark; refer Table 4 for details. For speculative decoding methods, one generation step consists of a draft forward pass followed by a verifier forward pass. We report the number of tokens decoded per step, τ , which comprises one draft and one verify step. However, this metric does not account for differences in drafter size and can therefore be misleading, particularly because the baseline drafters are smaller than ours. We consequently measure throughput using the 1K/8K throughput test described earlier. Following standard practice, we do not report accuracy for lossless methods because any differences arise from sampling randomness and numerical nondeterminism. We report τ under sampler configurations optimized for (1) system throughput and (2) per-request throughput. For each method, we select these configurations through a grid search over sampling hyperparameters; see Table 18. We use thinking mode for all methods. Although DFlash recommends disabling it for greater speedups (Z Lab, 2026), doing so reduces average accuracy from 76.36% to 55.40%; see Suppl. C.7. We therefore retain thinking mode for a fair comparison. When comparing with lossy methods, we instead report tokens per forward pass (TPF). For our method, TPF = τ /2 because the drafter and verifier are similar in size. For both drafting and verification, we use temperature temp = 1, top-p = 0.95, and top-k = 50 unless otherwise specified. We use a context length of 32,768, the native context length of Qwen3-8B, for all benchmarks and methods. 5.2.2

UnoQwen : Qwen-based Diffusion Augmented LLM (Ours)

Autoregressive Weights We initialize the AR weights from the open-source Qwen3-8B checkpoint (Yang et al., 2025) and keep them frozen during training. Training Diffusion Weights For each AR weight matrix, we add a rank-128 LoRA adapter with αLoRA = 256. These adapters introduce 0.35B trainable parameters. We train them for three epochs, corresponding to 14.7B tokens, on OpenThoughts3-1.2M (Guha et al., 2025), using a maximum sequence length of 4,096. We use a block-size curriculum with B ∈ {2, 4, 6, 8, 12, 16}, increasing the block size every half epoch. Training uses a global batch size of 64 and a constant learningrate of 10−5 with 2% warmup steps. We set the loss coefficients α = 0 and β = 1. Training takes approximately 32 hours on 4 nodes, each with 8 H200 GPUs. For ablations, we train all models for one epoch with B = 8, denoted Uno1ep Qwen . Unless stated otherwise, we evaluate this variant using the Linear sampler with B = 16. 12

5.2.3

Results

Comparison with Lossless Speculative Decoding Methods Table 2 compares the average number of toQwen kens per step, τ , for UnoQwen , EAGLE-3, and DFlash at temp = 1. For each method, we identify the configurations that maximize system and per-request throughput through a grid search over sampling parameters, as detailed in Table 18. At the largest supported batch size, with Linear sampler B = 4, all methods achieve their highest system throughput. UnoQwen exceeds 5700 tokens per second, outperforming DFlash and EAGLE-3 and achieving a 1.6× speedup over the Per-request Throughput (toks / sec) base AR model. This improvement enables faster RL training and offers practical speedups in LLM infer- Figure 2: System versus per-request throughput ence. For the best per-request throughput (batch size across batch sizes (concurrency C). Uno Paretoof 1), UnoQwen and EAGLE-3 perform best with tree- dominates speculative decoding and achieves up to 2.5× speedup over the base AR model. Parenthesampler configurations (B, K, V ) = (16, 32, 32) and ses indicate the sampler configuration yielding the (8, 32, 60), respectively, while DFlash performs best highest throughput for Uno. Exact throughputs are with B = 16. UnoQwen substantially outperforms both in Table 18. baselines, achieving a 2.5× speedup over the base AR model. It is also Pareto-dominant in throughput across all batch sizes Fig. 2. Unlike our method, which shares a single KV cache, EAGLE-3 and DFlash maintain separate caches for drafting and verification, resulting in higher peak memory usage. (B4)

System Throughput (toks / sec)

(B8)

(B8)

C64

(B8)

C32

(B16)

C16

(B16,V32) (B16,V32)

C8

C4

C2 C1

Takeaway 4: Uno is strictly faster than DFlash and EAGLE-3 across all batch sizes, with fewer additional parameters and a shared draft-verifier KV cache that reduces peak memory. Comparison with (Lossy) Diffusion Methods Table 3 compares UnoQwen with the leading lossy diffusion-based acceleration methods for completeness. Unlike our method UnoQwen , I-DLM, TiDAR, Jacobi Forcing, FLARE, Fast-dLLM-v2, and SDAR achieve lossy speedups relative to their respective parent models. Despite optimizing for quality first, our method achieves a higher TPF than most methods across numerous benchmarks. Benchmarks highlighted in red indicate accuracy degradation. FLARE supports both lossy and target-lossless sampling, but the paper does not state which produced the results. 5.2.4

Ablation

We ablate the main training choices in Table 12 using Uno1ep Qwen with the Linear sampler and B = 16. Ablation 1: Loss Terms The diffusion loss in (3) combines the distillation loss, LDCD , and the total variation loss, LTV . Training with LTV alone achieves a TPF of 2.39, outperforming both LDCD + LTV (2.23) and LDCD alone (2.23). Further analysis revealed that the LDCD loss was an order of magnitude larger than the LTV loss. Consequently, reducing the LDCD weight to 0.01 yielded a marginal improvement, increasing the TPF to 2.40. Ablation 2: Training Curriculum In Suppl. C.6.2, we show that increasing the training block size from 4 to 16 every half epoch improves TPF from 2.65 to 2.71, compared with using a fixed block size of 16 for two epochs. Ablation 3: Diffusion Weights Configuration In Suppl. C.6.3, we show that distributing diffusion adapters across all layers is more effective than placing the same number of parameters in only a subset of layers. We also ablate the LoRA rank rLoRA and scale αLoRA . Increasing rLoRA from 128 to 256 improves TPF, as shown in Table 12, but also increases inference cost. Across the tested values of αLoRA /rLoRA , a ratio of 64 performs best. 13

Table 2: Acceptance lengths (τ ), throughput (1K/8K test), peak memory usage, and additional parameter counts for UnoQwen , EAGLE-3, and DFlash at sampling temp = 1. We report the τ values that maximize system and per-request throughput. System Throughput Optimal UnoQwen EAGLE-3 DFlash B=4 B=4 B=4

Per-request Throughput Optimal UnoQwen EAGLE-3 DFlash B=16, V =32 B=8, V =60 B=16

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26

3.99 4.11 4.08 4.11 4.08

2.26 2.14 2.11 2.13 2.13

2.47 2.28 1.96 1.90 1.93

6.29 6.89 6.83 6.81 6.71

3.83 3.52 3.41 3.46 3.49

3.38 3.24 2.76 2.56 2.57

Coding HumanEval MBPP LCBv6

3.83 3.87 3.77

2.12 2.16 2.04

2.35 2.22 1.76

5.63 5.87 5.34

3.71 3.65 3.45

3.03 3.04 2.22

Science and Knowledge GPQA GPQA-Diamond MMLU-Pro

3.79 3.78 3.84

1.99 1.98 2.03

1.98 1.92 2.11

5.49 5.49 5.70

3.25 3.19 3.31

2.57 2.50 2.80

Instruction Following IFEval

3.48

1.91

1.92

4.58

3.49

2.26

τ Throughput (Toks / sec; ↑) Peak Memory (GiB; ↓) Additional Params (B; ↓)

3.89 5733 122.2 0.35

2.08 4944 130.0 0.40

2.07 5351 130.1 1.05

5.97 445 118.0 0.35

3.48 284 129.4 0.40

2.74 370 129.8 1.05

6

Related Work

Speculative Decoding Speculative decoding methods such as EAGLE-3 and DFlash also provide lossless speedups but require training a separate draft model that is smaller than the target model. Designing a draft model requires numerous choices about its depth, hidden dimension, MLP expansion, attention heads, parameter sharing, and KV projections for target features. Li et al. (2026a); Sandler et al. (2026) on the other hand use large pretrained dLLMs for drafting and separate AR models for verification. These large drafters incur substantial memory and latency overhead, making these methods slower than DFlash. In contrast, Uno uses a single architecture with distinct drafting and verification pathways. Consequently, Uno maintains a single KV cache, introduces fewer additional parameters, and requires less peak inference memory than conventional speculative decoding. These advantages yield larger speedups at high batch sizes, making Uno a drop-in replacement for methods that use separate draft and target models. Self-Speculative Decoding Sahoo et al. (2025b) introduced a hybrid AR-diffusion framework in which a single denoising model performs either multi-token diffusion generation or AR generation, depending on the available compute budget. This approach substantially outperforms block diffusion (Arriola et al., 2025). Building on self-speculative decoding (Stern et al., 2018), Liu et al. (2025) extended this framework by using the same denoising model to draft sequences in diffusion mode and verify them via rejection sampling in AR mode, yielding significant qualitative improvements over standalone diffusion generation. However, because this method modifies the base AR model’s weights to support diffusion generation, it does not preserve the model’s original distribution. Moreover, its speedups are limited to small batch sizes (Fu et al., 2026). In contrast, Uno preserves the base AR model’s distribution while providing lossless speedups, making it a drop-in alternative to these methods. Diffusion Models Large-scale d-LLMs (Nie et al., 2025; Gat et al., 2025; Bie et al., 2026; Fu et al., 2026; DiffusionGemmaTeam et al., 2026) can generate faster than similarly sized AR models 14

Table 3: Benchmark accuracy (ACC) and TPF for our lossless method, UnoQwen , and lossy diffusion methods. Entries are reported as ACCTPF . For UnoQwen , TPF uses temp = 0 and the tree sampler with (B, K, V ) = (16, 32, 60). Lossy speedups are taken from the respective papers, except for Fast-dLLM v2 and SDAR, which we evaluated using the provided checkpoints. Accuracy drops relative to the corresponding parent AR model are shaded in red by severity. Jacobi ForcingJ uses separate math and coding models trained from Qwen2.5-Math-7B-Instruct and Qwen2.5-Coder-7BInstruct, respectively. F Fast-dLLM v2 uses Qwen2.5-7B-Instruct. Lossy Speedup Methods Benchmark Model size Parent AR

UnoQwen 8B

Math GSM8K 96.13.56 MATH500 96.43.92 AIME-24 76.74.01 AIME-25 76.73.96 AIME-26 73.34.05 Coding HumanEval 94.83.67 89.04.21 MBPP LCBv6 51.43.75 Science and Knowledge GPQA 58.04.10 GPQA-D 60.94.01 MMLU-Pro 74.83.61 Instruction Following 86.52.49 IFEval

Jacobi FastFLARE LLaDA2.1TiDAR SDAR (Trust Diff) OPDLM I-DLM Forcing (MR)J dLLM v2F (AR Trust)L Flash (S Mode) 8B 8B 8B 8B 7B 7B 9B 100B-A5B Qwen3-8B Qwen2.5-7B Family Qwen3.5-9B N/A 91.42.3 72.02.8 10.02.8 10.02.6 –

80.47.1 – – – –

87.11 71.21 14.71 12.41 –

95.0N/A 96.8N/A 69.6N/A 60.8N/A –

91.44.0 – – – –

83.73.0 61.12.1 6.672.7 0.02.6 –

93.3N/A 95.2N/A 63.3N/A 54.4N/A –

– – – 63.35.4 –

75.62.8 68.11.5 16.6

57.97.3 65.410.0 –

59.81 48.71 9.71

93.32.6 92.2N/A 45.7N/A

83.54.1 70.42.8 –

63.42.5 63.04.5 10.01.7

92.1N/A 91.1N/A 49.7N/A

– – 44.16.5

– 40.2 56.91.6

– – –

– 36.11 53.71

54.9N/A 55.6N/A 73.1N/A

– – –

31.9N/A 21.72.2 –

– 71.2N/A 77.4N/A

– 66.74.0 75.34.4

61.41.5

–

50.11

84.7N/A

–

61.41.5

71.4N/A

83.42.2

at small batch sizes but generally lag behind them in quality. Their lossy speedups also diminish at larger batch sizes, limiting their applicability. Post-training these models is more expensive because rollouts become slower at large batch sizes and existing RL recipes and objectives require nontrivial modifications (Zhao et al., 2026; Wang et al., 2025). In contrast, Uno provides lossless speedups across batch sizes, uses standard AR post-training algorithms without modification, and substantially accelerates post-training. Multi-Token Prediction (MTP) Methods such as Medusa (Cai et al., 2024) and the approach of Gloeckle et al. (2024) provide lossless speedups but modify the transformer architecture to predict future tokens. Li et al. (2026b) showed that speculative decoding with a separately trained draft model is faster than these approaches. We further show that Uno outperforms speculative decoding methods and, consequently, these MTP methods. Nevertheless, architectural changes for predicting future tokens are complementary to Uno and could improve its draft acceptance rate. We leave this combination to future work. Quadratic Samplers Each generation step in our sampler uses two forward passes: one for drafting and one for verification. Quadratic sampling, proposed by Samragh et al. (2025) and used in TiDAR and Nemotron-Labs-Diffusion, can combine these operations into a single pass. For a draft of k tokens, it inserts k masked placeholders after each draft position, yielding k 2 masked positions that represent possible future continuations. This structure verifies the current draft while generating candidates for the next step. Practical speedups require efficient kernels, which we leave to future work. Comparison with Concurrent work Concurrent methods such as OPDLM, FLARE, and I-DLM provide lossy speedups relative to the base AR model because they fine-tune its weights; see Table 3. Although I-DLM is lossy, its I-DLM (R-ISD) variant uses gated LoRA. Despite the authors’ theoretical argument that its sampler is lossless, we could not reproduce lossless speedups using the released LoRA adapters and sampler implementation. Because this method is compatible with our proposed Ψ-spec sampler using a mask prior, we evaluate the combination and show that it achieves lossless speedups, although it remains substantially slower than our method. We provide a detailed comparison with I-DLM, including an evaluation of its official LoRA checkpoint, in Suppl. C.5. 15

7

Conclusion

In this work, we introduced diffusion-augmented LLMs, which unify two distinct generation pathways within a single architecture: an autoregressive pathway parameterized by θAR and a diffusion pathway parameterized by θAR +θ∆ . We showed that training the diffusion weights requires orders of magnitude fewer tokens than the AR weights. This formulation allows Uno either to be trained from scratch or to be created by augmenting an existing AR model with lightweight diffusion weights for faster generation. Uno serves as a drop-in alternative to speculative diffusion methods such as EAGLE-3 and DFlash, without requiring a separate draft model, and to self-speculative methods such as TiDAR, without introducing lossy speedups. Across the evaluated batch sizes, Uno achieves up to a 3× speedup over the base AR model, while retaining up to a 2× speedup at the largest batch size supported by that model. These gains accelerate both inference and RL post-training, where rollout generation is a major computational bottleneck. More broadly, diffusion-augmented LLMs expand the design space of AR language models by combining the quality of AR generation with the parallelism of diffusion in a unified architecture. Acknowledgement We thank A. Feder Cooper, Willis Guo, Rupesh Srivastava, and Matthew Yang for their helpful feedback and insightful discussions.

References Amazon Web Services. AGENTPERF01-BP01: Define performance-aligned success criteria for agent workloads. AWS Well-Architected Framework, Agentic AI Lens, 2026. URL https://docs.aws.amazon.com/wellarchitected/latest/agentic-ai-lens/ agentperf01-bp01.html. Marianne Arriola, Subham Sekhar Sahoo, Aaron Gokaslan, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Justin T Chiu, and Volodymyr Kuleshov. Block diffusion: Interpolating between autoregressive and diffusion language models. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=tyEyYT267x. Artificial Analysis. Announcing artificial analysis long context reasoning (AA-LCR), 2025. URL https://artificialanalysis.ai/articles/announcing-aa-lcr. Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17981–17993, 2021a. Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc V. Le, and Charles Sutton. Program synthesis with large language models, 2021b. URL https://arxiv.org/abs/2108.07732. Mislav Balunović, Jasper Dekoninck, Ivo Petrov, Nikola Jovanović, and Martin Vechev. MathArena: Evaluating llms on uncontaminated math competitions. In Advances in Neural Information Processing Systems, 2025. URL https://arxiv.org/abs/2505.23281. Victor Barres, Honghua Dong, Soham Ray, Xujie Si, and Karthik Narasimhan. τ 2 -bench: Evaluating conversational agents in a dual-control environment, 2025. URL https://arxiv.org/abs/ 2506.07982. Douglas Biber. A corpus-driven approach to formulaic language in English: Multi-word patterns in speech and writing. International Journal of Corpus Linguistics, 14(3):275–311, 2009. doi: 10.1075/ijcl.14.3.08bib. URL https://doi.org/10.1075/ijcl.14.3.08bib. 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, 16

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. URL https: //arxiv.org/abs/2602.08676. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020. 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. In Fortyfirst International Conference on Machine Learning, 2024. URL https://openreview.net/ forum?id=PEpbUobfJv. Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023. Jian Chen, Yesheng Liang, and Zhijian Liu. DFlash: Block diffusion for flash speculative decoding, 2026. URL https://arxiv.org/abs/2602.06036. Mark Chen et al. Evaluating large language models trained on code, 2021. URL https://arxiv. org/abs/2107.03374. Shuang Cheng, Yihan Bian, Dawei Liu, Yuhua Jiang, Yihao Liu, Linfeng Zhang, Qian Yao, Zhongbo Tian, Wenhai Wang, Qipeng Guo, et al. Sdar: A synergistic diffusion-autoregression paradigm for scalable sequence generation. In Findings of the Association for Computational Linguistics: ACL 2026, pp. 22058–22075, 2026. Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finetuned language models. Journal of machine learning research, 25(70):1–53, 2024. Kenneth Ward Church and Patrick Hanks. Word association norms, mutual information, and lexicography. Computational Linguistics, 16(1):22–29, 1990. URL https://aclanthology.org/ J90-1003/. 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. URL https://arxiv.org/ abs/2110.14168. Jasper Dekoninck, Nikola Jovanović, Tim Gehrunger, Kári Rögnvalddson, Ivo Petrov, Chenhao Sun, and Martin Vechev. Beyond benchmarks: MathArena as an evaluation platform for mathematics with llms, 2026. URL https://arxiv.org/abs/2605.00674. Justin Deschenaux, Caglar Gulcehre, and Subham Sekhar Sahoo. The diffusion duality, chapter II: /psi-samplers and efficient curriculum. In The Fourteenth International Conference on Learning Representations, 2026. URL https://openreview.net/forum?id=RSIoYWIzaP. DiffusionGemmaTeam, Adrien Ali TaÃŕga, James Assiene, Daniele Calandriello, Rahma Chaabouni, JoÃčo Gante, Tamara von Glehn, Nate Keating, Chris Knutsen, Martin Kukla, Tianlin Liu, Ivan Lobov, Ofir Nabati, JoÃčo Gabriel Oliveira, Nicolas Perez-Nieves, Nastasia Prutianova, Bobak Shahriari, Jean Tarbouriech, Pavel Tyletski, ÃĞaħlar ÃIJnlÃij, Cindy Wu, Glenn Cameron, Jerome Connor, Sertan Girgin, Maarten Grootendorst, Alon Levkovitch, Eliya Nachmani, Omar Sanseviero, Piotr Stanczyk, Quentin Berthet, Andrew Campbell, ClÃl’ment Crepy, Valentin De Bortoli, Arnaud Doucet, Romuald Elie, Alexandre Galashov, Klaus Greff, Alexis Jacq, David Ruhe, Yu-Han Wu, Sebastian Flennerhag, Brendan O’Donoghue, George Scrivener, and Shantanu Thakoor. Diffusiongemma technical report, 2026. URL https://arxiv.org/abs/2608.00146. Stefano Ermon. Introducing Mercury 2. Inception, 2026. URL https://www.inceptionlabs. ai/blog/introducing-mercury-2. 17

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. Technical report, NVIDIA, 2026. Technical report. Itai Gat, Heli Ben-Hamu, Marton Havasi, Daniel Haziza, Jeremy Reizenstein, Gabriel Synnaeve, David Lopez-Paz, Brian Karrer, and Yaron Lipman. Set block decoding is a language model inference accelerator. arXiv preprint arXiv:2509.04185, 2025. GLM-5-Team, :, Aohan Zeng, Xin Lv, Zhenyu Hou, Zhengxiao Du, Qinkai Zheng, Bin Chen, Da Yin, Chendi Ge, Chenghua Huang, Chengxing Xie, Chenzheng Zhu, Congfeng Yin, Cunxiang Wang, Gengzheng Pan, Hao Zeng, Haoke Zhang, Haoran Wang, Huilong Chen, Jiajie Zhang, Jian Jiao, Jiaqi Guo, Jingsen Wang, Jingzhao Du, Jinzhu Wu, Kedong Wang, Lei Li, Lin Fan, Lucen Zhong, Mingdao Liu, Mingming Zhao, Pengfan Du, Qian Dong, Rui Lu, Shuang-Li, Shulin Cao, Song Liu, Ting Jiang, Xiaodong Chen, Xiaohan Zhang, Xuancheng Huang, Xuezhen Dong, Yabo Xu, Yao Wei, Yifan An, Yilin Niu, Yitong Zhu, Yuanhao Wen, Yukuo Cen, Yushi Bai, Zhongpei Qiao, Zihan Wang, Zikang Wang, Zilin Zhu, Ziqiang Liu, Zixuan Li, Bojie Wang, Bosi Wen, Can Huang, Changpeng Cai, Chao Yu, Chen Li, Chengwei Hu, Chenhui Zhang, Dan Zhang, Daoyan Lin, Dayong Yang, Di Wang, Ding Ai, Erle Zhu, Fangzhou Yi, Feiyu Chen, Guohong Wen, Hailong Sun, Haisha Zhao, Haiyi Hu, Hanchen Zhang, Hanrui Liu, Hanyu Zhang, Hao Peng, Hao Tai, Haobo Zhang, He Liu, Hongwei Wang, Hongxi Yan, Hongyu Ge, Huan Liu, Huanpeng Chu, Jia’ni Zhao, Jiachen Wang, Jiajing Zhao, Jiamin Ren, Jiapeng Wang, Jiaxin Zhang, Jiayi Gui, Jiayue Zhao, Jijie Li, Jing An, Jing Li, Jingwei Yuan, Jinhua Du, Jinxin Liu, Junkai Zhi, Junwen Duan, Kaiyue Zhou, Kangjian Wei, Ke Wang, Keyun Luo, Laiqiang Zhang, Leigang Sha, Liang Xu, Lindong Wu, Lintao Ding, Lu Chen, Minghao Li, Nianyi Lin, Pan Ta, Qiang Zou, Rongjun Song, Ruiqi Yang, Shangqing Tu, Shangtong Yang, Shaoxiang Wu, Shengyan Zhang, Shijie Li, Shuang Li, Shuyi Fan, Wei Qin, Wei Tian, Weining Zhang, Wenbo Yu, Wenjie Liang, Xiang Kuang, Xiangmeng Cheng, Xiangyang Li, Xiaoquan Yan, Xiaowei Hu, Xiaoying Ling, Xing Fan, Xingye Xia, Xinyuan Zhang, Xinze Zhang, Xirui Pan, Xu Zou, Xunkai Zhang, Yadi Liu, Yandong Wu, Yanfu Li, Yidong Wang, Yifan Zhu, Yijun Tan, Yilin Zhou, Yiming Pan, Ying Zhang, Yinpei Su, Yipeng Geng, Yong Yan, Yonglin Tan, Yuean Bi, Yuhan Shen, Yuhao Yang, Yujiang Li, Yunan Liu, Yunqing Wang, Yuntao Li, Yurong Wu, Yutao Zhang, Yuxi Duan, Yuxuan Zhang, Zezhen Liu, Zhengtao Jiang, Zhenhe Yan, Zheyu Zhang, Zhixiang Wei, Zhuo Chen, Zhuoer Feng, Zijun Yao, Ziwei Chai, Ziyuan Wang, Zuzhou Zhang, Bin Xu, Minlie Huang, Hongning Wang, Juanzi Li, Yuxiao Dong, and Jie Tang. Glm-5: from vibe coding to agentic engineering, 2026. URL https://arxiv.org/abs/2602.15763. Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Synnaeve. Better & faster large language models via multi-token prediction. arXiv preprint arXiv:2404.19737, 2024. Google Cloud. The roi of ai 2025: Measuring the impact of ai and ai agents. Technical report, Google Cloud, 2025. URL https://cloud.google.com/resources/content/ roi-of-ai-2025. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023. Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 1474–1487, 2020. Etash Guha, Ryan Marten, Sedrick Keh, Negin Raoof, Georgios Smyrnis, Hritik Bansal, Marianna Nezhurina, Jean Mercat, Trung Vu, Zayne Sprague, Ashima Suvarna, Benjamin Feuer, Liangyu Chen, Zaid Khan, Eric Frankel, Sachin Grover, Caroline Choi, Niklas Muennighoff, Shiye Su, Wanjia Zhao, John Yang, Shreyas Pimpalgaonkar, Kartik Sharma, Charlie Cheng-Jie Ji, Yichuan 18

Deng, Sarah Pratt, Vivek Ramanujan, Jon Saad-Falcon, Jeffrey Li, Achal Dave, Alon Albalak, Kushal Arora, Blake Wulfe, Chinmay Hegde, Greg Durrett, Sewoong Oh, Mohit Bansal, Saadia Gabriel, Aditya Grover, Kai-Wei Chang, Vaishaal Shankar, Aaron Gokaslan, Mike A. Merrill, Tatsunori Hashimoto, Yejin Choi, Jenia Jitsev, Reinhard Heckel, Maheswaran Sathiamoorthy, Alexandros G. Dimakis, and Ludwig Schmidt. Openthoughts: Data recipes for reasoning models, 2025. URL https://arxiv.org/abs/2506.04178. Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025. Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022. Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Liang Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. Iclr, 1(2):3, 2022. Lanxiang Hu, Siqi Kou, Yichao Fu, Samyam Rajbhandari, Tajana Rosing, Yuxiong He, Zhijie Deng, and Hao Zhang. Fast and accurate causal parallel decoding using jacobi forcing. arXiv preprint arXiv:2512.14681, 2025. Declan Jackson, William Keating, George Cameron, and Micah Hill-Smith. AA-Omniscience: Evaluating cross-domain knowledge reliability in large language models, 2025. URL https: //arxiv.org/abs/2511.13029. 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. In International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=chfJJYC3iL. Carlos E. Jimenez, John Yang, Alexander Wettig, Shunyu Yao, Kexin Pei, Ofir Press, and Karthik Narasimhan. SWE-bench: Can language models resolve real-world github issues? In International Conference on Learning Representations, 2024. URL https://arxiv.org/abs/2310. 06770. Minseo Kim, Minjae Lee, Seunghyuk Oh, Kevin Galim, Donghoon Kim, Coleman Hooper, Harman Singh, Amir Gholami, Hyung Il Koo, and Wonjun Kang. Efficientrollout: System-aware selfspeculative decoding for rl rollouts. arXiv preprint arXiv:2606.18967, 2026. Kimi KimiTeam, Tongtong Bai, Yifan Bai, Yiping Bao, Jianfeng Cai, Xinyuan Cai, Peizhou Cao, Yuxuan Cao, Ziwei Chai, Y Charles, et al. Kimi k3: Open frontier intelligence. arXiv preprint arXiv:2607.24653, 2026. Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696–21707, 2021. Seul Lee, Karsten Kreis, Srimukh Prasad Veccham, Meng Liu, Danny Reidenbach, Yuxing Peng, Saee Paliwal, Weili Nie, and Arash Vahdat. Genmol: A drug discovery generalist with discrete diffusion. arXiv preprint arXiv:2501.06158, 2025. Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pp. 19274–19286. PMLR, 2023. Guanghao Li, Zhihui Fu, Min Fang, Qibin Zhao, Ming Tang, Chun Yuan, and Jun Wang. Diffuspec: Unlocking diffusion language models for speculative decoding. In Findings of the Association for Computational Linguistics: ACL 2026, pp. 20896–20910, 2026a. Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle-3: Scaling up inference acceleration of large language models via training-time test. Advances in Neural Information Processing Systems, 38:136737–136756, 2026b. 19

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. In International Conference on Learning Representations, 2024. URL https://arxiv.org/abs/2305.20050. 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. arXiv preprint arXiv:2511.08923, 2025. Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.16834, 2024. 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. arXiv preprint arXiv:2502.09992, 2025. Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. In The Thirteenth International Conference on Learning Representations, 2025. URL https: //openreview.net/forum?id=sMyXP8Tanm. Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35: 27730–27744, 2022. Long Phan, Alice Gatti, Ziwen Han, Nathaniel Li, Josephina Hu, Hugh Zhang, Chen Bo Calvin Zhang, Mohamed Shaaban, John Ling, Sean Shi, et al. Humanity’s last exam. arXiv preprint arXiv:2501.14249, 2025. David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R. Bowman. GPQA: A graduate-level google-proof q&a benchmark. In First Conference on Language Modeling, 2024. URL https://openreview. net/forum?id=Ti67584b98. Subham S. Sahoo. Foundations of Diffusion Language Models. PhD thesis, Cornell University, 2026. URL https://www.proquest.com/dissertations-theses/ foundations-diffusion-language-models/docview/3355013570/se-2. Copyright Database copyright ProQuest LLC; ProQuest does not claim copyright in the individual underlying works; Last updated - 2026-06-23. Subham Sekhar Sahoo, Marianne Arriola, Aaron Gokaslan, Edgar Mariano Marroquin, Alexander M Rush, Yair Schiff, Justin T Chiu, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024a. URL https://openreview.net/forum?id=L4uaAR4ArM. Subham Sekhar Sahoo, Aaron Gokaslan, Christopher De Sa, and Volodymyr Kuleshov. Diffusion models with learned adaptive noise. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b. URL https://openreview.net/forum?id=loMa99A4p8. Subham Sekhar Sahoo, Justin Deschenaux, Aaron Gokaslan, Guanghan Wang, Justin T Chiu, and Volodymyr Kuleshov. The diffusion duality. In Forty-second International Conference on Machine Learning, 2025a. URL https://openreview.net/forum?id=9P9Y8FOSOk. Subham Sekhar Sahoo, Zhihan Yang, Yash Akhauri, Johnna Liu, Deepansha Singh, Zhoujun Cheng, Zhengzhong Liu, Eric Xing, John Thickstun, and Arash Vahdat. Esoteric language models: A family of any-order diffusion llms. arXiv preprint arXiv:2506.01928, 2025b. Subham Sekhar Sahoo, Jean-Marie Lemercier, Zhihan Yang, Justin Deschenaux, Jingyu Liu, John Thickstun, and Ante Jukic. Scaling beyond masked diffusion language models. arXiv preprint arXiv:2602.15014, 2026. 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. arXiv preprint arXiv:2507.11851, 2025. 20

Jameson Sandler, Jacob Christopher, Tom Hartvigsen, and Ferdinando Fioretto. Specdiff-2: Scaling diffusion drafter alignment for faster speculative decoding. Proceedings of Machine Learning and Systems, 8:1128–1147, 2026. Yair Schiff, Subham Sekhar Sahoo, Hao Phung, Guanghan Wang, Sam Boshar, Hugo Dalla-torre, Bernardo P de Almeida, Alexander M Rush, Thomas PIERROT, and Volodymyr Kuleshov. Simple guidance mechanisms for discrete diffusion models. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id= i5MrJ6g5G1. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024. Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias. Simplified and generalized masked diffusion for discrete data. Advances in neural information processing systems, 37: 103131–103167, 2024. Quan Shi, Alexandra Zytek, Pedram Razavi, Karthik Narasimhan, and Victor Barres. τ -knowledge: Evaluating conversational agents over unstructured knowledge. arXiv preprint arXiv:2603.04370, 2026. Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pp. 2256–2265. PMLR, 2015. Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020. Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 32211–32252. PMLR, 23–29 Jul 2023. URL https://proceedings.mlr.press/v202/song23a.html. Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models. In Neural Information Processing Systems, 2018. URL https://api. semanticscholar.org/CorpusID:53208380. Xingyu Su, Jacob Helwig, Shubham Parashar, Atharv Chagi, Lakshmi Jotsna, Degui Zhi, James Caverlee, Dileep Kalathil, and Shuiwang Ji. Data-efficient autoregressive-to-diffusion language models via on-policy distillation. arXiv preprint arXiv:2606.06712, 2026. Sophia Tang, Yinuo Zhang, and Pranam Chatterjee. Peptune: De novo generation of therapeutic peptides with multi-objective-guided discrete diffusion. ArXiv, pp. arXiv–2412, 2025. The Terminal-Bench Team. Terminal-bench 2.1, may 2026. URL https://www.tbench.ai/ news/terminal-bench-2-1. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. Shashank Verma and Neal Vaidya. Mastering LLM techniques: Inference optimization. NVIDIA Technical Blog, November 2023. URL https://developer.nvidia.com/blog/ mastering-llm-techniques-inference-optimization/. Accessed 2026-08-24. Guanghan Wang, Gilad Turok, Yair Schiff, Marianne Arriola, and Volodymyr Kuleshov. d2: Improving reasoning in diffusion language models via trajectory likelihood estimation. arXiv preprint arXiv:2509.21474, 2025. 21

Yubo Wang, Xueguang Ma, Ge Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, Aaran Arulraj, Xuan He, Ziyan Jiang, Tianle Li, Max Ku, Kai Wang, Alex Zhuang, Rongqi Fan, Xiang Yue, and Wenhu Chen. MMLU-Pro: A more robust and challenging multi-task language understanding benchmark. In Advances in Neural Information Processing Systems, 2024. URL https://arxiv.org/abs/2406.01574. Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le. Finetuned language models are zero-shot learners. arXiv preprint arXiv:2109.01652, 2021. 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. URL https://arxiv.org/abs/2509.26328. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jing Zhou, Jingren Zhou, Junyang Lin, Kai Dang, Keqin Bao, Kexin Yang, Le Yu, Lianghao Deng, Mei Li, Mingfeng Xue, Mingze Li, Pei Zhang, Peng Wang, Qin Zhu, Rui Men, Ruize Gao, Shixuan Liu, Shuang Luo, Tianhao Li, Tianyi Tang, Wenbiao Yin, Xingzhang Ren, Xinyu Wang, Xinyu Zhang, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yinger Zhang, Yu Wan, Yuqiong Liu, Zekun Wang, Zeyu Cui, Zhenru Zhang, Zhipeng Zhou, and Zihan Qiu. Qwen3 technical report, 2025. URL https://arxiv.org/abs/2505.09388. Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, et al. Dapo: An open-source llm reinforcement learning system at scale. Advances in Neural Information Processing Systems, 38:113222–113244, 2026a. Xingkai Yu. Nano-vLLM: A lightweight vLLM implementation built from scratch. https:// github.com/GeeeekExplorer/nano-vllm, 2025. Software repository. Yifan Yu, Yuqing Jian, Junxiong Wang, Zhongzhu Zhou, Donglin Zhuang, Xinyu Fang, Sri Yanamandra, Xiaoxia Wu, Qingyang Wu, Shuaiwen Leon Song, et al. Introspective diffusion language models. arXiv preprint arXiv:2604.11035, 2026b. Xiangchi Yuan, Dachuan Shi, Chunhui Zhang, Zheyuan Liu, Shenglong Yao, Soroush Vosoughi, and Wenke Lee. Behavior knowledge merge in reinforced agentic models. In Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 33007–33028, 2026. Z Lab. Qwen3-8B-DFlash-b16 model card. Qwen3-8B-DFlash-b16, 2026. Accessed 2026-08-03.

https://huggingface.co/z-lab/

Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. d1: Scaling reasoning in diffusion large language models via reinforcement learning. Advances in Neural Information Processing Systems, 38:56729–56762, 2026. Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E. Gonzalez, Clark Barrett, and Ying Sheng. SGLang: Efficient execution of structured language model programs, 2024. URL https://arxiv.org/ abs/2312.07104. Jeffrey Zhou, Tianjian Lu, Swaroop Mishra, Siddhartha Brahma, Sujoy Basu, Yi Luan, Denny Zhou, and Le Hou. Instruction-following evaluation for large language models, 2023. URL https: //arxiv.org/abs/2311.07911. Yuchen Zhu, Jing Shi, Chongjian Ge, Hao Tan, Yiran Xu, Wanrong Zhu, Jason Kuen, Koustava Goswami, Rajiv Jain, Yongxin Chen, et al. Flare: Diffusion for hybrid language model. arXiv preprint arXiv:2606.01774, 2026.

22

Contents 1

Introduction

2

2

Background

3

2.1

Autoregressive Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

3

2.2

Discrete Diffusion Language Models . . . . . . . . . . . . . . . . . . . . . . . . . . . .

3

3

4

5

Diffusion-augmented LLMs

4

3.1

Autoregressive and Diffusion Pathways . . . . . . . . . . . . . . . . . . . . . . . . . .

4

3.2

Diffusion Distillation Phase . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

5

3.3

Practical Considerations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

6

Ψ-Speculative Sampler

6

4.1

Diffusion Sampler . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

7

4.2

One-step Diffusion w/ AR Verification . . . . . . . . . . . . . . . . . . . . . . . . . . .

7

4.3

Inference-Time Scaling . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

8

Experiments

8

5.1

End to End Diffusion Augmented LLM Training . . . . . . . . . . . . . . . . . . . . .

9

5.2

Open Weights LLMs . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11

6

Related Work

14

7

Conclusion

16

A Background

24

A.1 Masked Diffusion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 A.2 Uniform-State Diffusion Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 A.3 Self-speculative Decoding . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25 A.4 Ψ-Samplers . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25 A.5 Discrete Consistency Distillation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 26 B Ψ-Spec samplers

27

B.1 Diffusion Proposal Distribution Extended . . . . . . . . . . . . . . . . . . . . . . . . . 27 B.2 Single-Step Generation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27 B.3 Sampling Algorithm . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27 C Additional Experiments

28

C.1 Experiment Configs . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 28 C.2 Uno RL . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 28 C.3 Uno Evaluations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 29 C.4 Qwen Finetuned on OpenThoughts . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 30 23

C.5 UnoQwen vs I-DLM (R-ISD) Comparison . . . . . . . . . . . . . . . . . . . . . . . . . 30 C.6 UnoQwen Ablations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 32 C.7 DFlash Thinking Mode Ablation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 36 C.8 Extended Comparison vs EAGLE-3 and DFlash . . . . . . . . . . . . . . . . . . . . . 36

A

Background

A.1

Masked Diffusion

MDMs use a masked prior, where π = m ∈ V is the one-hot representation of a special [MASK] token. During the forward process (1), tokens either remain unchanged or transition to the masked state m, after which they stay masked. This behavior carries over to the reverse process. The posMDM terior of the reverse process qs∣t for 0 ≤ s < t < 1 can be derived using Bayes’ Rule, and is given by: ⎧ s ℓ s −αt ⎪ xℓ + 1−α z) ⎪ Cat ( ⋅ ; α1−α ℓ 1−αt t t [qs∣t (⋅ ∣ zt , x)]MDM = ⎨ ℓ ⎪ Cat( ⋅ ; x ) ⎪ ⎩

if zℓt = m, otherwise.

(6)

The learned reverse posterior is MDM (zs ∣ zt , x = xθ (zt , t)), [pθs∣t (zs ∣ zt )]MDM = qs∣t

where xθ ∶ V L × [0, 1] → ∆L is the denoising model. A key limitation is that once unmasked, tokens cannot be remasked. This can create compounding errors during inference, as the denoising model xθ imperfectly models the clean data. NELBO The likelihood of a sequence under the learned reverse diffusion process is generally intractable to evaluate exactly, as it requires marginalizing over all possible unmasking trajectories. Therefore, typically the negative evidence lower bound (NELBO) is minimized, which is a tractable upper bound on the negative log-likelihood. For continuous-time masked diffusion, the NELBO simplifies to the weighted denoising objective (Sahoo et al., 2024a; Shi et al., 2024; Ou et al., 2025) LMDM NELBO = Et∼U [0,1],qt

A.2

αt′ log⟨xℓθ (zt , t), xℓ ⟩. 1 − αt

Uniform-State Diffusion Models

We now focus on Uniform-State Diffusion Models (USDMs), which use a uniform prior π = 1/K (Austin et al., 2021a; Sahoo et al., 2025a). Unlike masked diffusion, where tokens are fixed once unmasked, USDMs allow continuous updates of all token positions, which naturally supports self-correction at inference. This capability makes them particularly effective for few-step generation (Sahoo et al., 2025a), inference-time scaling (Deschenaux et al., 2026) and guided sampling (Schiff et al., 2025). Given a noisy sequence zt at time t, the Markov property defines the reverse posterior qs∣t for obtaining a less noisy sequence zs (0 ≤ s < t) by sampling from the following distribution: ℓ [qs∣t (⋅ ∣ zt , x)]USDM = Cat (⋅ ;

Kαt zℓt ⊙ xℓ + (αt∣s − αt )zℓt + (αs − αt )xℓ + (1 − αt∣s )(1 − αs )1/K ) . (7) Kαt ⟨zℓt , xℓ ⟩ + 1 − αt

Similar to the posterior for masked diffusion models above, the true reverse posterior (7) depends on the unknown clean sequence x, hence it is approximated as: USDM [pθs∣t (zs ∣ zt )]USDM = qs∣t (zs ∣ zt , x = xθ (zt , t)),

where xθ ∶ V L × [0, 1] → ∆L is the denoising network parameterized by θ. 24

NELBO

This posterior induces the following NELBO (Sahoo et al., 2025a): ℓ ℓ ℓ ℓ LUSDM NELBO (q, pθ ; x) = −Et∼U [0,1],qt ∑ f(zt , xθ (zt , t), αt ; x ), ℓ∈[L]

where f(zℓt , xℓθ (zℓt , t),αt ; xℓ ) = −K

⎡ (x̄ℓ ) αt ′ ⎢ ⎢ K − K − (ζt 1 ℓ ℓ + 1 ℓ ℓ ) ∑ log θ r ⎢ ℓ zt =x zt ≠x ℓ Kαt ⎢ x̄r (x̄θ )r (x̄ℓθ )j j ⎣

⎤ ⎥ (x̄ℓ )r αt 1 log θℓ 1zℓ ≠xℓ − ((K − 1)ζt 1zℓ =xℓ − 1zℓ ≠xℓ ) log ζt ⎥ ⎥. t 1 − αt ζt t ⎥ (x̄θ )i t ⎦

Here, x̄ℓ = Kαt xℓ + (1 − αt )1, x̄ℓθ = Kαt xℓθ (zt , t) + (1 − αt )1, αt′ denotes the time derivative of t αt , r = arg maxj∈[K] (zℓt )j is the nonzero entry of zt , ζt = Kα1−α , and i denotes the index in x t +1−αt corresponding to 1, that is, xi = 1.

A.3

Self-speculative Decoding

Self-speculative decoding eliminates the need for a separate draft model by using a single network for both proposal and verification. Stern et al. (2018) first proposed augmenting the base model with auxiliary prediction heads to predict a block of future tokens in parallel, then using the same model to verify these predictions and commit the longest validated prefix. Their exact-match verification rule guarantees equivalence to greedy decoding, rather than exact sampling. Liu et al. (2025) instead convert an AR model into a diffusion model: the model drafts tokens in diffusion mode and verifies them in AR mode, with the attention pattern changing between the two modes while the model weights remain the same in both modes. They combine this approach with the rejection-sampling acceptance rule of Leviathan et al. (2023) and Chen et al. (2023) to recover exact sampling from the base model’s distribution. Concurrent work (Zhu et al., 2026; Yu et al., 2026b) builds on the same idea. Crucially, however, this approach requires modifying the base AR model’s weights, so it is not lossless.

A.4

Ψ-Samplers

Ψ-Samplers (Deschenaux et al., 2026) are predictor-corrector samplers for discrete diffusion models that share the marginals of the ancestral sampler. They interpolate between the reverse posterior and the forward process (1): Ψs∣t (⋅ ∣ x, zt ; π) = κt qs∣t (⋅ ∣ zt , x; π) + (1 − κt )qs (⋅ ∣ x; π), where π is the noise prior (e.g., the masked prior π = m for MDMs or the uniform prior π = 1/K for USDMs), qs∣t (⋅ ∣ ⋅ ; π) is the reverse posterior for MDMs (6) or USDMs (7), qs is the forward noising process, and κt ∈ [0, 1] sets the correction strength. The first term predicts by denoising the current state; the second corrects by re-injecting noise, so earlier decisions can be revised. Setting κt = 1 recovers ancestral sampling. Since x is unavailable at generation time, the practical sampler substitutes the denoiser prediction xθ (zt , t), giving the token-wise transition [Ψθs∣t (⋅ ∣ zt ; xθ (zt , t), π)]ℓ ℓ ℓ = κt qs∣t (⋅ ∣ zt , xθ (zt , t); π) + (1 − κt ) [αs q0∣t (⋅ ∣ zt , xθ (zt , t); π) + (1 − αs )π] .

Ancestral samplers can lock in incorrect early predictions, whereas Ψ-samplers correct errors explicitly: for MDMs the corrector can remask decoded tokens, and for USDMs it assigns every token non-zero probability, permitting transitions into states that the denoiser incorrectly scores as lowprobability. Sample quality therefore keeps improving with additional steps. For MDMs, the transition (2) simplifies to [Ψθs∣t (⋅ ∣ zt ; xθ (zt , t), π = m)]ℓ ⎧ ℓ ⎪ ⎪ ⎪ Cat(⋅ ; [αs + κt (1 − αs )]xθ (zt , t) + (1 − κt )(1 − αs )m) , =⎨ κt αt (1−αs ) t (1−αs ) ⎪ ⎪ ] xℓθ (zt , t) + [1 − αs + κt α1−α ] m) , ⎪ Cat(⋅ ; [αs − 1−αt t ⎩ 25

zℓt ≠ m, zℓt = m.

(8)

To obtain the marginals for single-step generation, we simply plug in s = 0 and t = 1 into (8) which gives: [Ψθs=0∣t=1 (⋅ ∣ zt ; xθ (zt , t), π = m)]ℓ ⎧ ℓ ⎪ ⎪ ⎪ Cat(⋅ ; [1 + κt (1 − 1)] xθ (z1 , 1) + (1 − κt )(1 − 1)m) , =⎨ ⎪ ⎪ Cat(⋅ ; [αs − κt ⋅0⋅(1−1) ] xℓθ (z1 , 1) + [1 − 1 + κt ⋅0⋅(1−1) ] m) , ⎪ 1−0 1−0 ⎩ ℓ ℓ ⎧ ⎪ ⎪ Cat(⋅; xθ (z1 )), z1 ≠ m, =⎨ ℓ ℓ ⎪ ⎪ ⎩ Cat(⋅; xθ (z1 )), z1 = m. = Cat(⋅; xℓθ (z1 ))

zℓ1 ≠ m, zℓ1 = m.

(9)

For USDM, the transition (2) simplifies to: [Ψθs∣t (⋅ ∣ zt ; xθ (zt , t), π = 1/K)]ℓ ⎡ ⎢ ⎛ 1 ⎢Kαt [αs + κt (1 − αs )] zℓ ⊙ [xθ (zt , t)]ℓ = Cat ⋅ ; t ℓ ℓ ⎝ Kαt ⟨zt , [xθ (zt , t)] ⟩ + 1 − αt ⎢⎢ ⎣ + κt (1 − αs )αt∣s zℓt + [αs (1 − αt ) − κt αt (1 − αs )] [xθ (zt , t)]ℓ ⎤ 1 ⎥⎞ + (1 − αs ) [κt (1 − αt∣s ) + (1 − κt ) (Kαt ⟨zℓt , [xθ (zt , t)]ℓ ⟩ + 1 − αt )] ⎥⎥ . K ⎥⎠ ⎦ To obtain the marginals for single-step generation, we simply plug in s = 0 and t = 1 into (7) which gives [Ψθs=0∣t=1 (⋅ ∣ zt ; xθ (zt ), π = 1/K)]ℓ = Cat(⋅; xℓθ (z1 , 1)).

A.5

(10)

Discrete Consistency Distillation

Sahoo et al. (2025a) show that Uniform-state discrete diffusion emerges from an underlying Gaussian diffusion process (Sohl-Dickstein et al., 2015; Song et al., 2020; Kingma et al., 2021) defined on the one-hot representation xℓ ∈ V: applying an arg max projection to the Gaussian latents recovers the discrete process. This correspondence lets us transfer techniques from the continuous to the discrete domain, most notably distillation for few-step generation. Distillation requires deterministic trajectories, namely the probability-flow ODE, which exist for Gaussian but not for discrete diffusion. Since discrete latents are simply arg max projections of Gaussian ones, Sahoo et al. (2025a) construct the trajectory in continuous space and map it back to the discrete domain. Discrete Consistency Distillation (DCD) trains few-step generators on these trajectories. First, the Gaussian probability-flow ODE is constructed on the one-hot vectors and the resulting continuous states are mapped back to tokens via arg max, yielding a deterministic sequence of discrete latents. The teacher xθ0 ∶ V L × [0, 1] → ∆L is a neural network with parameters θ0 that are held fixed, while the student xθ ∶ V L × [0, 1] → ∆L is a neural network whose parameters θ are trained. Given two adjacent states along this path, a noisier zt at time t and a less noisy zs at time s = t − T , the student takes the noisier state as input and is trained to match the teacher’s clean-token distribution at the less noisy one, in the spirit of consistency distillation (Song et al., 2023; Sahoo et al., 2025a): L

ℓ 1∶L LDCD (θ; θ0 ) = ∑ DKL (xℓθ (z1∶L t , t) ∥ xθ0 (zt , s)) , ℓ=1

Note that the teacher’s parameters θ0 are not updated during training. Distillation proceeds over numerous rounds, with the step size ∆ increased at every round; at the end of each round the trained student parameters θ are copied into the teacher, which then serves as the target for the next round. By taking progressively larger denoising steps while remaining consistent along a shared trajectory, DCD compresses a many-step USDM sampler into a few-step one, making Uniform-state diffusion particularly effective in the low-step regime. 26

Limitation: DCD is trained on deterministic probability-flow trajectories, which differ from the stochastic denoising trajectories followed at sampling time. This train-test mismatch limits its effectiveness.

B

Ψ-Spec samplers

B.1

Diffusion Proposal Distribution Extended

Given a noisy block zt , we use the Ψ-Spec transition in (2) to sample a less noisy block zs , where s < t. Because the denoiser uses next-token-prediction (NTP) parameterization, the clean-token L distributions used to construct zs are given by x1∶B−1 θAR ,θ∆ ([x , zt ]). The resulting transition is obtained by drawing each token in parallel from its corresponding marginal distribution: ℓ 1∶B−1 ⎧ ⎪ ⎪ Ψs∣t (⋅ ∣ zt ; xθAR ,θ∆ (⋅), π) , zℓs ∼ ⎨ ℓ=1 1∶B−1 ⎪ ⎪ ⎩ Ψs∣t (⋅ ∣ zt ; xθAR (⋅), π) .

ℓ>1

(11)

ℓ=1

(12)

Notice that the logits for the first, clean position use only the base AR parameters θAR , as shown in (11), whereas the logits at the noisy positions use both the AR weights and the diffusion adapters. This separation prevents distribution shift because the diffusion adapters are trained only to denoise noisy tokens. We compute (11) and (12) in a single forward pass using the gated LoRA technique of Samragh et al. (2025).

B.2

Single-Step Generation

After removing time conditioning from the denoising model and shifting the logits one position to the left to account for its next-token prediction (NTP) formulation in (9) for MDMs and (10) for USDMs, we show that the resulting marginal distribution has the following form. Let Ψ0 denote the distribution induced by the Ψ-sampler at t = 0: Ψℓ0 = {

B.3

xℓθAR ,θ∆ ([xL , zt ]), ℓ > 1 xℓθAR ([xL , zt ]). ℓ=1

Sampling Algorithm

Algo. 1 shows Uno’s complete sampling algorithm for the single-sequence case. Differences to speculative decoding (Leviathan et al., 2023) are written in orange. Algorithm 1 Lossless Ψ-Spec decoding for USDMs (Linear Sampler) Require: Prefix x of length L, target model xθAR ,θ∆ , block size B i.i.d.

∼ ∏B−1 U[V] ▷ Initialize a block of B − 1 random tokens 1 2: [q0 , q] ← xθAR ,θ∆ ([xL , z]; x<L ) ▷ θAR apply on xL ; θAR + θ∆ on z L+1 3: Sample x ∼ q0 ▷ free clean token from draft step 4: Sample x̃ ∼ q ▷ B − 1 tokens sampled in parallel 5: p ← xθAR ([xL+1 , x̃]; x) ▷ AR verifier distribution defined by θAR i.i.d. B−1 6: sample r ∼ ∏1 U[0, 1] p 7: n ← min({i ∈ [B − 1] ∶ ri > min(1, qi )} ∪ {B}) i 8: if n = B then 9: sample z̃ ∼ pB 10: else 11: sample z̃ ∼ Norm([pn − qn ]+ ) 12: end if 13: return [x, xL+1 , x̃1∶n−1 , z̃] 1: Sample noise z

27

C

Additional Experiments

C.1

Experiment Configs

C.1.1

Benchmarks

We provide additional details about the set of benchmarks we evaluate Uno on in Table 4. Table 4: Benchmarks used in our Uno evaluation. “Generations” denotes the number of responses sampled per example when computing average pass@1. Benchmark

Generations

Agentic τ 2 -Bench

3

Description Multi-turn tool-use tasks in which an agent interacts with a simulated user and domain APIs while following domain-specific policies. Realistic and complex tasks that autonomous agents complete in command-line container environments. Human-validated software-engineering tasks requiring agents to resolve real GitHub issues by modifying their associated repositories.

Terminal-Bench v2.1

3

SWE-bench Verified

3

Long-Context Reasoning AA-LCR

3

Open-answer questions requiring information extraction, synthesis, and reasoning over long documents such as reports and legal texts.

Science and Knowledge AA-Omniscience

1

Humanity’s Last Exam

1

GPQA-Diamond

5

Closed-book questions measuring factual recall and calibration across a broad range of economically relevant domains. Difficult expert-level questions spanning mathematics, the sciences, the humanities, and other academic disciplines. Expert-validated, graduate-level multiple-choice questions in biology, physics, and chemistry.

Math GSM8K

2

MATH500

2

AIME 2024–2026

10

Coding MBPP

2

HumanEval Instruction Following IFEval

C.1.2

2

2

Grade-school mathematics word problems requiring multi-step arithmetic reasoning. A 500-problem subset of MATH covering competition-level mathematical reasoning with free-form answers. Integer-answer problems from the 2024, 2025, and 2026 American Invitational Mathematics Examinations. Short Python program-synthesis problems specified through naturallanguage descriptions and test cases. Hand-written Python function-completion problems evaluated using executable unit tests. Verifiable instruction-following tasks that test compliance with explicit, objectively checkable constraints.

UnoQwen , EAGLE-3, DFlash Sampler Configs for Qwen3-8B

At temp = 1, we use top-p = 0.95 and top-k = 50; temp = 0 denotes greedy decoding. UnoQwen uses linear blocks of size B ∈ {4, 8, 16} or, for tree verification, B = 16, tree size V = 60, and candidate top-K = 32. DFlash uses blocks B ∈ {4, 8, 16} and we evaluate both thinking-enabled and thinking-disabled variants (see Suppl. C.7). Linear EAGLE-3 B ∈ {4, 8, 16} draft tokens; its tree configuration uses seven draft steps, top-k = 10, and at most V = 60 draft tokens.

C.2

Uno RL

Starting from the supervised fine-tuning (SFT) checkpoint, four specialized experts were trained for mathematics, code generation, tool use, and browse search using the DAPO reinforcement learning algorithm (Yu et al., 2026a). The Math expert was trained for 2,560 steps and generated 28.9B response tokens, while the Code expert was trained in two stages of 600 and 1,500 steps, producing 28

Table 5: Open-source model checkpoints used in our evaluations. Model

Checkpoint

DiffusionGemma Nemotron-Labs-Diffusion EAGLE-3 DFlash Fast-dLLM v2 SDAR

https://huggingface.co/google/diffusiongemma-26B-A4B-it https://huggingface.co/nvidia/Nemotron-Labs-Diffusion-14B https://huggingface.co/AngelSlim/Qwen3-8B_eagle3 https://huggingface.co/z-lab/Qwen3-8B-DFlash-b16 https://huggingface.co/Efficient-Large-Model/Fast_dLLM_v2_7B https://huggingface.co/JetLM/SDAR-8B-Chat-b16

6B and 11.94B response tokens, respectively. Both experts used a learning rate of 5 × 10−7 , 16 groups, and 32 prompts per step, yielding 512 sequences per step across 32 nodes per expert. Math training used maximum context lengths of 32,768 and 65,536 tokens, while Code training supported contexts up to 65,536 and 128K tokens. The Tool-use expert completed 39 steps with 128 prompts per step and generated 1.4B tokens, including tool-call tokens. The Browse Search expert completed 59 steps with 64 prompts per step and generated 8.4B tokens, including tool-call tokens. For both Tool-use and Browse Search, the maximum trajectory length was 128K tokens, covering modelgenerated tokens as well as tool outputs. Finally, the four experts were consolidated into a single model using ISO-Merger (RAM; Yuan et al., 2026), a data-free approach that combines specialists trained from a shared base checkpoint without requiring additional rollouts, gradient updates, or distillation.

C.3

Uno Evaluations

Table 6: We report TPFs for a subset of the evaluations using different sampler configurations for the Uno model. Among all configurations, the Linear sampler with B = 4 achieves the highest system throughput at a batch size of 64, while the Tree sampler with (B, K, V ) = (16, 32, 32) achieves the highest per-user throughput (batch size 1). Dataset

Linear B4

Linear B8

Linear B16

Tree B16,K32,V32

Tree B16,K64,V64

Tree B16,K64,V32

GSM8K AIME-24 AIME-25 AIME-26 MATH500 HumanEval AA-LCR

1.9 1.8 1.8 1.8 1.8 1.8 1.8

2.3 2.1 2.2 2.1 2.1 2.6 2.1

2.4 2.2 2.3 2.1 2.2 2.8 2.2

2.7 2.5 2.4 2.4 2.4 3.4 2.6

2.6 2.5 2.6 2.5 2.5 2.5 2.5

2.8 2.4 2.6 2.4 2.5 3.4 2.4

TPF Highest Sys. Throughput Highest Per-req. Throughput

1.8 5190 275

2.2 5034 319

2.3 3445 345

2.6 2665 379

2.5 1573 339

2.6 2709 371

Table 7: DG, Uno, NLD, Mercury2. ∗ As reported by Artificial Analysis’ live tracker on August 30, 2026. Method

Max. Sys. Throughput (toks / sec)

Max. Per-req. Throughput (toks / sec)

Uno (Ours) AR (Ours) DiffusionGemma Nemotron-Labs-Diffusion Mercury-2

5255 3577 1136 2794 1197∗

383 176 836 290 769∗

29

Table 8: TPFs for the SFT checkpoint and its RL-post-trained counterpart, evaluated using the diffusion weights trained for the SFT checkpoint. Even after extensive RL post-training, the SFT diffusion adapters retained their speedup, with only a 6% reduction in TPFs.

C.4

SFT

Post-Trained

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26

2.66 2.27 2.17 2.26 2.22

1.95 1.93 1.93 1.97 1.83

Coding HumanEval MBPP

1.97 2.35

2.34 2.19

Science and Knowledge GPQA-Diamond 2.03

1.97

Instruction Following IFEval 2.12

2.50

Other Benchmarks HLE AA-Omniscience

2.14 2.57

2.15 2.34

TPF

2.25

2.10

Qwen Finetuned on OpenThoughts

Table 9 compares Qwen3-8B with a variant fine-tuned on OpenThoughts across a subset of benchmarks. Fine-tuning on OpenThoughts substantially degrades performance, reducing accuracy by up to 15% points, depending on the benchmark. Nevertheless, diffusion weights trained on OpenThoughts successfully accelerate generation despite the distribution mismatch between OpenThoughts and the data used to train the AR weights.

Table 9: This table compares Qwen3-8B with a variant fine-tuned on OpenThoughts across a subset of benchmarks. Qwen3-8B Qwen-finetuned

C.5

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26

96±0.2 96.2±0.3 77.7±3.0 70.7±4.2 67.7±3.9

95.2±1.1 94.8±0.0 68.3±17.6 53.3±0.0 61.7±5.9

Coding HumanEval MBPP LCBv6

94.4±1.0 88.7±0.4 50.9±1.2

77.8±4.6 78.1±0.4 42.5±0.8

UnoQwen vs I-DLM (R-ISD) Comparison

In this section, we compare I-DLM (R-ISD), the lossless variant of I-DLM, with our method, UnoQwen . In the first experiment, we train I-DLM (R-ISD) with a rank-128 LoRA adapter for one epoch on the OpenThoughts dataset, using exactly the same training configuration as Uno1ep Qwen . For inference, 30

we use the sampler provided in the official I-DLM repository3 . All results for I-DLM (R-ISD) in Table 10 are evaluated with a block size of B = 16. The results show that the I-DLM sampler degrades accuracy on numerous benchmarks relative to the Qwen AR model, as highlighted in red, whereas Uno1ep Qwen does not. Upon inspecting the I-DLM codebase, we find that its sampler performs greedy drafting4 , but its rejection-sampling verification procedure is not adjusted accordingly. As a result, the sampler does not preserve the claimed losslessness property. I-DLM (R-ISD) uses LoRA adapters and retains the transformer’s causal attention pattern. It is therefore compatible with our Ψ-Spec sampler using the masked prior π = m. In the second set of experiments, we evaluate I-DLM (R-ISD) with Ψ-Spec. As shown in Table 10, this combination remains lossless. We observe the same result when evaluating the I-DLM (R-ISD) adapters released in the official repository. However, both variants provide considerably lower speedups than our method.

Table 10: I-DLM and Uno1ep Qwen evaluation with the linear sampler at B = 16 and temp = 1. We compare the I-DLM checkpoints (Ours and Official) under Ψ-spec samplers with π = m; accuracy is reported with 95% confidence intervals across seeds, using two seeds for MMLU-Pro and ten seeds for all other benchmarks. Ψ-spec sampler

I-DLM sampler

Sampler Qwen3 AR

I-DLM (R-ISD) (ours)

I-DLM (R-ISD) (ours)

I-DLM (R-ISD) (official)

Acc. (%)

Acc. (%)

TPF

Acc. (%)

TPF

Acc. (%)

TPF

Acc. (%)

TPF

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26

96.0±0.2 96.2±0.3 77.7±3.0 70.7±4.2 67.7±3.9

96.0±0.2 96.3±0.5 77.3±2.9 63.0±3.6 68.7±4.7

2.14 2.36 2.33 2.38 2.34

96.0±0.2 96.1±0.4 76.7±4.2 68.7±4.8 64.3±3.2

1.95 2.06 2.03 2.06 2.02

96.0±0.2 95.9±0.3 76.7±2.8 70.7±5.1 66.7±3.2

2.23 2.06 1.80 1.79 1.78

96.0±0.2 96.2±0.3 77.7±2.8 69.7±3.2 67.0±3.3

2.66 2.97 2.85 2.92 2.84

Coding HumanEval MBPP LCBv6

94.4±1.0 88.7±0.4 50.9±1.2

93.7±0.9 87.8±0.4 48.9±1.5

2.30 2.19 2.05

95.0±0.6 88.7±0.7 50.4±1.4

1.77 1.80 1.68

94.5±0.5 88.7±0.5 51.0±1.4

1.98 1.93 1.60

94.3±0.7 88.8±0.4 50.1±1.4

2.40 2.47 2.30

Science and Knowledge GPQA GPQA-Diamond MMLU-Pro

56.8±0.7 59.5±0.9 74.7±0.7

57.2±1.1 58.4±0.5 74.6±0.5

2.01 2.02 2.07

56.1±1.1 59.4±1.4 74.7±0.2

1.66 1.67 1.75

56.6±1.1 60.6±1.2 75.1±0.9

1.71 1.69 1.84

57.0±1.0 59.7±1.1 74.7±0.7

2.35 2.35 2.46

Instruction Following IFEval

85.8±0.7

85.3±0.7

1.87

85.9±0.8

1.65

85.7±0.6

1.85

86.0±0.7

1.96

Average

76.6±0.4

75.6±0.6

2.17

76.0±0.7

1.84

76.5±0.5

1.86

76.4±0.4

2.55

Model

Uno (ours)

Additionally, Table 11 provides TPF for the I-DLM LoRA checkpoint evaluated on block sizes B = 4, B = 8, and B = 16. The results show a significant increase in TPF from B = 4 to B = 8, while TPF seems to saturate beyond B = 8.

3 https://github.com/Introspective-Diffusion/I-DLM/commit/a23c1a12ef997c7f3ad616b25bcfb62db39ded68 4 https://github.com/Introspective-Diffusion/I-DLM/blob/a23c1a12ef997c7f3ad616b25bcfb62db39ded68/

inference/sglang/sglang/srt/dllm/algorithm/idlm_blockN.py#L575

31

Table 11: TPFs for the official I-DLM (R-ISD) model with our Ψ-spec samplers across block sizes B. B=4 B=8 B=16 Math GSM8K MATH500 AIME-24 AIME-25 AIME-26

1.87 1.80 1.68 1.66 1.67

2.20 2.04 1.80 1.81 1.78

2.23 2.06 1.80 1.79 1.78

Coding HumanEval MBPP LCBv6

1.75 1.73 1.54

1.96 1.93 1.60

1.98 1.93 1.60

Science and Knowledge GPQA GPQA-Diamond MMLU-Pro

1.61 1.59 1.69

1.71 1.69 1.83

1.71 1.69 1.84

Instruction Following IFEval

1.73

1.86

1.85

Average

1.69

1.85

1.86

C.6

UnoQwen Ablations

C.6.1

LoRA Rank and Loss Ablations

In Table 12, we ablate the LoRA rank of Uno1ep Qwen , as well as various weightings of KL and TV loss terms. We see that average TPF increases from 2.39 to 2.47 when increasing the LoRA rank from 128 to 256, while the number of trainable parameters increases from 349M to 698M. It turns out that a loss weighting of 0.01×KL plus 1×TV slightly outperforms a pure TV loss objective. Note that a factor of 0.01 for the KL term amounts to a KL loss term only about 10× smaller in magnitude than the TV term, as the KL loss is naturally around 10× larger than the TV loss. C.6.2

Training Block Size Curriculum

Here we ablate two block-size curricula for UnoQwen training. Table 13 shows the TPF attained by two different three-epoch block-size curricula, both initialized from a model trained on block size 2 for 0.5 epochs, and block size 4 for another 0.5 epochs (written as 0.5@2, 0.5@4). The first curriculum is our default curriculum used for UnoQwen , which continues at 0.5@6, 0.5@8, 0.5@12, and 0.5@16; the second one continues at 2@16. Average TPF drops from 2.71 to 2.65 when employing the second curriculum, which shows that incrementally increasing block size during training is beneficial over consistently training at a large block size.

32

Table 12: TPF results for one-epoch Uno1ep Qwen checkpoints trained with TV-only, KL-only, and weighted KL+TV objectives. The TV-only objective is evaluated using rank-128 LoRA adapters (the default UnoQwen configuration, highlighted in blue) and rank-256 adapters, while all other objectives use rank-128 adapters. All adapters are applied to every projection and use αLoRA /r = 2. All checkpoints are evaluated using the linear sampler with B = 16 and temp = 1. TV only r=128 r=256 α=0, β=1

Loss ablations (r=128) KL only α=1, β=0

KL + TV α=1, β=1

0.1 KL + TV α=0.1, β=1

0.05 KL + TV α=0.05, β=1

0.01 KL + TV α=0.01, β=1

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26

2.47 2.73 2.65 2.76 2.81

2.53 2.82 2.75 2.80 2.82

2.27 2.52 2.47 2.55 2.53

2.30 2.56 2.51 2.56 2.51

2.39 2.65 2.62 2.66 2.59

2.42 2.70 2.66 2.72 2.67

2.46 2.71 2.70 2.75 2.68

Coding HumanEval MBPP LCBv6

2.26 2.33 2.17

2.37 2.40 2.22

2.04 2.13 1.94

2.07 2.14 1.97

2.30 2.26 2.08

2.23 2.29 2.13

2.25 2.32 2.17

Science and Knowledge GPQA GPQA-Diamond MMLU-Pro

2.21 2.22 2.31

2.27 2.26 2.36

1.99 2.00 2.11

2.03 2.03 2.14

2.15 2.15 2.24

2.19 2.18 2.28

2.20 2.21 2.30

Instruction Following IFEval

1.79

1.99

2.08

1.93

2.08

2.08

2.02

Average TPF

2.39

2.47

2.22

2.23

2.35

2.38

2.40

Table 13: TPF for two three-epoch TV-only block curricula. The notation e@B denotes e epochs of training at block size B. Both models first train on 0.5@2, 0.5@4. The standard curriculum then progressively increases the block size, whereas the fixed-B = 16 curriculum trains for the remaining two epochs entirely at B = 16. Models are evaluated with the linear sampler at B = 16 and temp = 1. 0.5@2, 0.5@4, 0.5@6, 0.5@8, 0.5@12, 0.5@16

0.5@2, 0.5@4, 2@16

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26

2.83 3.13 3.18 3.14 3.03

2.81 3.10 3.04 3.00 3.02

Coding HumanEval MBPP LCBv6

2.58 2.60 2.37

2.49 2.58 2.37

Science and Knowledge GPQA GPQA-Diamond MMLU-Pro

2.45 2.47 2.57

2.45 2.46 2.56

Instruction Following IFEval

2.20

1.95

Average

2.71

2.65

33

C.6.3

LoRA Adapter Ablations

We ablate which projection matrices in the transformer layer to apply LoRA weights to in Table 14. By default, we apply LoRA to all projections (this includes the Q, K, V, O matrices from the attention layer, as well as the MLP gate, up, and down projections); in addition, we ablate only applying LoRA to the attention layers (i.e., Q, K, V, and O projections); only to the Q and K projections of the attention layer; only to the Q projections; and only to the O projections, which is the setup recommended in Fu et al. (2026). In each setting, we adapted the rank to keep parameter parity at 349M LoRA parameters. All models for this ablation were trained on one epoch at block size eight. Table 14 shows that at fixed αLoRA / rank = 2, applying LoRA to all projections work best. Consequently, we apply LoRA all to projections for our Uno models. Table 14: TPF for Uno1ep Qwen LoRA target-projection ablations with αLoRA /rLoRA = 2. Ranks are chosen to approximately maintain LoRA parameter parity. All checkpoints are evaluated with the linear sampler at B=16 and temp = 1. All projections

Q, K, V, O projections

Q, V projections

Q projection

O projection

Rank r LoRA αLoRA αLoRA /r

128 256 2

364 728 2

729 1458 2

1184 2368 2

1184 2368 2

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26

2.47 2.73 2.65 2.76 2.81

2.45 2.68 2.67 2.70 2.64

2.32 2.55 2.48 2.54 2.46

2.20 2.40 2.35 2.39 2.32

2.43 2.69 2.72 2.71 2.68

Coding HumanEval MBPP LCBv6

2.26 2.33 2.17

2.23 2.31 2.15

2.16 2.22 2.05

2.11 2.11 1.95

2.33 2.33 2.15

Science and Knowledge GPQA GPQA-Diamond MMLU-Pro

2.21 2.22 2.31

2.19 2.21 2.29

2.10 2.11 2.20

2.02 2.01 2.09

2.20 2.21 2.30

Instruction Following IFEval

1.79

2.00

1.88

1.73

1.78

Average

2.39

2.38

2.26

2.14

2.38

C.6.4

Varying LoRA α/r

For LoRA with rank rLoRA and a projection Wbase ∈ Rn×m of the base model, both LoRA training and inference use αLoRA W = Wbase + BA, rLoRA where B ∈ Rn×r , A ∈ Rr×m , and αLoRA > 0 is a hyperparameter that controls the overall contribution of the LoRA adapter. In Table 15 we ablate the ratio αLoRA /rLoRA across values between 2 and 256 or both one-epoch and three-epoch models. All three-epoch models in this ablation were trained on the block size curriculum described in Suppl. C.6.2. The table shows that the optimal ratio differs depending on the training horizon, and αLoRA /rLoRA = 16 is optimal at three epochs, while αLoRA /rLoRA = 64 is optimal at one epoch. Consequently, we pick αLoRA /rLoRA = 16 for UnoQwen , while for Uno, we set αLoRA /rLoRA = 64 which we observed to be preferable.

34

Table 15: TPF for one-epoch and three-epoch, TV-only Uno1ep Qwen models with rank-128 LoRA adapters applied to all projections, across LoRA scaling values. All checkpoints are evaluated with the linear sampler at B = 16 and temp = 1; averages are unweighted across the 12 benchmarks. (a) One epoch of training. αLoRA = 256

αLoRA = 512

αLoRA = 1024

αLoRA = 2048

αLoRA = 8192

αLoRA = 16384

αLoRA = 32768

αLoRA /r

2

4

8

16

64

128

256

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26

2.47 2.73 2.65 2.76 2.81

2.53 2.79 2.77 2.83 2.72

2.63 2.89 2.85 2.91 2.77

2.66 2.97 2.85 2.92 2.84

2.76 3.04 3.00 3.02 2.93

2.76 3.02 2.92 3.01 2.93

2.72 2.98 2.90 2.90 2.82

Coding HumanEval MBPP LCBv6

2.26 2.33 2.17

2.39 2.39 2.21

2.36 2.43 2.26

2.40 2.47 2.30

2.55 2.56 2.32

2.45 2.51 2.31

2.51 2.46 2.25

Science and Knowledge GPQA GPQA-Diamond MMLU-Pro

2.21 2.22 2.31

2.26 2.27 2.36

2.31 2.32 2.41

2.35 2.35 2.46

2.41 2.40 2.51

2.41 2.40 2.50

2.37 2.36 2.45

Instruction Following IFEval

1.79

2.01

2.13

1.96

2.01

2.27

1.84

Average

2.39

2.46

2.52

2.55

2.63

2.62

2.54

(b) Three epochs of training.

αLoRA = 2048

αLoRA = 8192

αLoRA = 16384

αLoRA /r

16

64

128

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26

2.83 3.13 3.18 3.14 3.03

2.83 3.12 3.02 3.08 3.05

2.78 3.06 2.94 3.00 2.97

Coding HumanEval MBPP LCBv6

2.58 2.60 2.37

2.49 2.56 2.34

2.44 2.51 2.30

Science and Knowledge GPQA GPQA-Diamond MMLU-Pro

2.45 2.47 2.57

2.43 2.47 2.55

2.41 2.39 2.50

Instruction Following IFEval

2.20

2.00

2.13

Average

2.71

2.66

2.62

35

C.7

DFlash Thinking Mode Ablation

In Table 16 we evaluate DFlash (Chen et al., 2026) with thinking mode enabled and disabled, at temp = 1. Our results show that while disabling thinking mode significantly increases tokens per step, it deteriorates accuracy from around 76% to roughly 55%. This quality difference does not contradict losslessness: enabling thinking changes the chat template and hence the target-model distribution that the speculative decoder preserves. In light of these results, we subsequently enable thinking mode for DFlash (Z Lab, 2026). Table 16: DFlash thinking ablation at temp = 1, top-p = 0.95, and top-k = 50. Accuracy in percent. AL denotes average acceptance length, i.e., tokens per verification step. DFlash, B=4 No thinking

DFlash, B=16

Thinking

No thinking

Thinking

Acc. (%)

AL

Acc. (%)

AL

Acc. (%)

AL

Acc. (%)

AL

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26

93.93 84.00 23.33 20.00 16.67

3.15 3.19 2.95 2.95 3.01

95.91 96.60 76.67 60.00 63.33

2.47 2.28 1.96 1.90 1.93

93.56 82.80 30.00 16.67 20.00

5.75 6.23 5.41 4.88 4.76

95.83 96.20 76.67 73.33 66.67

3.38 3.24 2.76 2.56 2.57

Coding HumanEval MBPP LCBv6

87.20 67.40 25.14

3.68 3.10 2.64

94.51 89.60 49.71

2.35 2.22 1.76

87.20 67.40 24.00

10.10 5.17 4.90

94.51 90.00 49.71

3.03 3.04 2.22

Science and Knowledge GPQA GPQA-Diamond MMLU-Pro

43.53 43.43 66.45

2.57 2.50 2.70

57.59 61.11 74.75

1.98 1.92 2.11

44.42 44.44 66.53

3.61 3.70 4.12

55.80 58.08 74.83

2.57 2.50 2.80

Instruction Following IFEval

85.21

1.84

86.14

1.92

87.80

2.59

84.66

2.26

Average

54.69

2.86

75.49

2.07

55.40

5.10

76.36

2.74

C.8

Extended Comparison vs EAGLE-3 and DFlash

By default, we evaluate DFlash with block sizes B ∈ {8, 16}. For EAGLE-3, we use a linear block size of B = 8, as well as its standard tree configuration with depth D = 7 (corresponding to block size B = 8), top-k = 10, and a verification budget of V = 60 tokens. These EAGLE-3 tree parameters follow the official settings provided in the paper and repository. We provide extended evaluation results for UnoQwen , EAGLE-3, and DFlash across various sampler settings in Table 17.

36

Table 17: Tokens per step across samplers for UnoQwen , EAGLE-3, and thinking-enabled DFlash. Values at both temp = 0 and temp = 1 are unfiltered. (a) UnoQwen linear and tree samplers.

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26 Coding HumanEval MBPP LCBv6 Science and Knowledge GPQA GPQA-Diamond MMLU-Pro Instruction Following IFEval Average

Linear B=4

Linear B=8

temp = 0 Linear Tree B=16 B=16, V =32

Tree B=16, V =60

Linear B=4

Linear B=8

temp = 1 Linear Tree B=16 B=16, V =32

4.14 4.27 4.26 4.25 4.20

5.52 5.96 5.89 5.77 6.10

6.28 7.09 6.95 6.73 7.36

7.02 7.67 7.83 7.67 7.82

7.27 8.01 8.55 8.36 8.61

3.99 4.11 4.08 4.11 4.08

5.14 5.46 5.49 5.45 5.31

5.67 6.26 6.35 6.28 6.05

6.29 6.89 6.83 6.81 6.71

6.49 7.12 6.81 7.07 6.86

4.31 4.42 4.29

6.26 6.55 6.03

7.70 8.34 7.09

8.74 9.16 8.16

9.28 9.45 8.20

3.83 3.87 3.77

4.67 4.86 4.51

5.16 5.20 4.74

5.63 5.87 5.34

5.78 5.96 5.53

4.35 4.37 4.19

6.07 6.08 5.67

7.15 7.30 6.57

8.16 8.46 7.48

8.52 8.64 7.69

3.79 3.78 3.84

4.64 4.64 4.75

4.91 4.95 5.14

5.49 5.49 5.70

5.65 5.65 5.90

4.20 4.27

5.57 5.96

6.83 7.11

7.29 7.95

7.85 8.37

3.48 3.89

4.36 4.94

4.40 5.42

4.58 5.97

4.53 6.11

Tree B=16, V =60

(b) EAGLE-3 linear and tree samplers.

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26 Coding HumanEval MBPP LCBv6 Science and Knowledge GPQA GPQA-Diamond MMLU-Pro Instruction Following IFEval τ

Linear B=4

Linear B=8

temp = 0 Tree Tree Linear B=8, B=8, B=16 V =32 V =60

2.43 2.26 2.23 2.23 2.28

2.68 2.46 2.42 2.41 2.49

2.73 2.40 2.49 2.54 2.58

3.82 3.50 3.41 3.51 3.63

4.07 3.80 3.62 3.73 3.75

3.78 3.49 3.48 3.51 3.63

4.24 3.80 3.70 3.72 3.71

2.26 2.14 2.11 2.13 2.13

2.48 2.27 2.25 2.27 2.26

2.49 2.28 2.23 2.26 2.27

3.62 3.27 3.22 3.31 3.28

3.83 3.51 3.46 3.49 3.51

3.57 3.29 3.25 3.28 3.28

3.83 3.52 3.41 3.46 3.49

2.87 2.82 2.64

3.44 3.52 3.08

3.70 3.58 3.04

4.52 4.80 4.27

5.07 5.02 4.47

4.94 4.97 4.41

5.21 5.34 4.65

2.12 2.16 2.04

2.36 2.31 2.14

2.28 2.30 2.15

3.39 3.39 3.22

3.61 3.65 3.38

3.43 3.40 3.22

3.71 3.65 3.45

2.48 2.49 2.35

2.77 2.77 2.60

2.98 2.93 2.62

3.85 3.82 3.68

4.14 4.17 3.93

4.04 3.87 3.69

4.20 4.09 3.95

1.99 1.98 2.03

2.09 2.07 2.14

2.10 2.07 2.14

3.06 3.00 3.11

3.25 3.19 3.27

3.05 2.99 3.10

3.25 3.19 3.31

2.10 2.43

3.06 2.81

3.39 2.91

3.66 3.87

3.29 4.09

4.32 4.01

4.60 4.27

1.91 2.08

2.22 2.24

2.14 2.23

3.16 3.25

3.04 3.43

3.07 3.24

3.49 3.48

Tree B=16, V =32

Tree B=16, V =60

Linear B=4

Linear B=8

temp = 1 Tree Tree Linear B=8, B=8, B=16 V =32 V =60

Tree B=16, V =32

Tree B=16, V =60

(c) Thinking-enabled DFlash at block sizes B=4, B=8, and B=16.

Math GSM8K MATH500 AIME-24 AIME-25 AIME-26 Coding HumanEval MBPP LCBv6 Science and Knowledge GPQA GPQA-Diamond MMLU-Pro Instruction Following IFEval τ

temp = 0 B=4 B=8 B=16

temp = 1 B=4 B=8 B=16

2.54 2.35 2.04 2.00 1.97

3.18 2.99 2.48 2.30 2.37

3.60 3.49 3.10 2.85 3.03

2.47 2.28 1.96 1.90 1.93

3.04 2.82 2.22 2.23 2.19

3.38 3.24 2.76 2.56 2.57

2.39 2.28 1.92

3.12 2.96 2.25

4.16 3.87 2.77

2.35 2.22 1.76

2.71 2.66 1.94

3.03 3.04 2.22

2.12 2.09 2.23

2.55 2.46 2.70

3.23 3.13 3.26

1.98 1.92 2.11

2.28 2.19 2.47

2.57 2.50 2.80

2.10 2.17

2.71 2.67

2.58 3.26

1.92 2.07

2.23 2.42

2.26 2.74

37

Table 18: Median per-request throughput (tok/s/stream) for the 1K-input/8K-output Qwen3-8B workload at temp = 1. Each value is the median of three fresh repetitions; “–” denotes a residentcapacity-infeasible point. Uno

EAGLE-3

Con.

B4

B8 B16 B16, V 32 B16, V 60

1 2 4 8 16 32 64 128

305 566 1086 1906 3247 4586 5733 –

366 703 1205 2397 3814 4873 5600 –

416 768 1420 2384 3436 4086 4112 –

445 803 1411 2161 2798 2958 3138 –

1 2 4 8 16 32 64 128

305 283 271 238 203 143 90 –

366 351 301 300 238 152 87 –

416 384 355 298 215 128 64 –

445 401 353 270 175 92 49 –

B4

DFlash

B8 B16 B16, V 32 B16, V 60

B4

AR

B8 B16

–

System throughput 435 204 198 161 730 401 377 312 1221 784 737 585 1630 1381 1292 1038 1869 2429 2233 1693 1964 3745 3308 2228 2004 4944 3965 2580 – – – –

289 531 906 1399 1746 1881 1969 –

284 287 331 369 176 504 538 645 702 352 771 1058 1188 1296 653 1019 1874 2073 2153 1144 1118 3075 3315 3007 1933 1171 4407 4246 3396 2796 – 5351 4900 3592 3662 – – – – –

Per-user throughput (tokens/s/user) 435 365 305 204 117 61 31 –

204 200 196 173 152 117 77 –

198 189 184 161 140 103 62 –

38

161 156 146 130 106 70 40 –

289 266 227 175 109 59 31 –

284 252 193 127 70 37 – –

287 269 264 234 192 138 84 –

331 322 297 259 207 133 77 –

369 351 324 269 188 106 56 –

176 176 163 143 121 87 57 –

Related documents

• Privacy-Preserving RAG by Concealing Sensitive Information from External LLMs#451247
• QuantGuard: Learnable Rounding for Repairing Quantization-Conditioned Backdoors in LLMs#677658
• QuantGuard: Learnable Rounding for Repairing Quantization-Conditioned Backdoors in LLMs#674457
• LLMs for Secure Hardware Design and Related Problems: Opportunities and Challenges#175114
• Generating Attacks for LLMs with GFlowNets#444582