Unlocking Lossless Speedups in LLMs via Discrete Diffusion

Summary: Large language models generate text one token at a time, which makes inference slow. Diffusion language models generate many tokens in parallel, but they lose quality, and their speedups fade when a server handles many requests at once. We introduce Uno: a standard autoregressive (AR) model plus small diffusion weights that draft a block of tokens in one forward pass, while the unchanged AR weights check the draft. The output follows exactly the AR model’s distribution. Uno is faster than leading speculative decoding methods at every batch size we tested, and can be built on top of existing open-weight models such as Qwen3-8B.

Large language models (LLMs) are trained to predict one token at a time. This works very well on a wide range of tasks, but it makes inference slow: the model produces only one token per forward pass. Long reasoning traces make this worse; so does reinforcement learning (RL) post-training, where most of the time goes into generating rollouts.

Sequential decoding also wastes hardware. During decoding, the GPU spends most of its time moving weights and the key-value (KV) cache from memory, and little time computing. Decoding is often memory-bound, especially at long context lengths: a forward pass that scores several tokens then costs about as much as a pass that scores one. A model that could produce several correct tokens per pass would improve GPU utilization and throughput (tokens/second) at the same time.

Two families of methods try to do this, and both have a catch.

• Speculative decoding (Leviathan et al., 2023) uses a small draft model to propose draft tokens, and the large target model (from which one wants to generate) checks them in one parallel pass. A rejection-sampling rule makes the output follow exactly the target model’s distribution, so the speedup is lossless. But you have to design, train, and serve a second model, with its own parameters and its own KV cache, which increases memory. Furthermore, the draft model itself generates tokens autoregressively.
• Diffusion language models (d-LLMs) generate a block of tokens in parallel by denoising it. Leading open d-LLMs are either trained from scratch (Nie et al., 2025) or obtained by fine-tuning an AR model (Fu et al., 2026; DiffusionGemma Team, 2026). Fine-tuning an AR model into a diffusion model has been shown to be cheaper and more effective than training from scratch (Gong et al., 2025; Ye et al., 2025). However, changing the AR model’s weights typically lowers its scores on benchmarks (DiffusionGemma Team, 2026). As a result, d-LLMs trade quality for speed, and even those speedups largely disappear at large batch sizes. In fact, prior work has shown that d-LLMs are slower than AR decoding at large batch sizes (Kim et al., 2025; DiffusionGemma Team, 2026).

The last point matters. In agentic workloads, one user request can start many parallel agents, tool calls, and retries, and servers batch these calls together; RL rollouts are batched too. Production cares about throughput at large batch sizes.

This raises the question: can we accelerate inference from an AR model by generating multiple tokens per forward pass, without changing the distribution tokens are sampled from? Our answer is a new model class that we call diffusion-augmented LLMs. We call our models Uno, because they unify AR weights and diffusion weights in one network.

One network, two pathways

A diffusion-augmented LLM uses a causal LLM backbone but carries two sets of weights in every layer.

• The AR weights $\theta_{\mathrm{AR}}$ are trained with standard LLM recipes (pre-training, supervised fine-tuning (SFT), and RL post-training).
• The diffusion weights $\theta_{\Delta}$ are LoRA adapters added to every linear layer of the AR model and trained in our proposed Diffusion Distillation phase, with $\theta_{\mathrm{AR}}$ kept frozen.

This gives two pathways that share the same base parameters. The AR pathway uses $\theta_{\mathrm{AR}}$ alone and performs sequential generation: given a clean sequence $\mathbf{x} \in \mathcal{V}^L$, where $\mathcal{V}$ denotes the vocabulary set, the causally masked output at position $\ell-1$ predicts the next token,

\[p_{\theta_{\mathrm{AR}}}(\mathbf{x}_{\ell} \mid \mathbf{x}_{<\ell}) = \mathbf{x}^{\ell-1}_{\theta_{\mathrm{AR}}}(\mathbf{x}), \qquad 1 \le \ell \le L.\]

The diffusion pathway uses $\theta_{\mathrm{AR}} + \theta_{\Delta}$ and performs a sequence-to-sequence map: given a noisy sequence $\mathbf{z}_t \in \mathcal{V}^L$, it predicts every clean token in parallel, keeping the next-token parameterization of the backbone,

\[p_{\theta_{\mathrm{AR}} + \theta_{\Delta}}(\mathbf{x}_{\ell} \mid \mathbf{z}_t) = \mathbf{x}^{\ell-1}_{\theta_{\mathrm{AR}} + \theta_{\Delta}}(\mathbf{z}_t), \qquad 1 \le \ell \le L.\]

The diffusion pathway drafts (blocks of) tokens, and the verification pathway checks them via rejection sampling against $p_{\theta_{\mathrm{AR}}}$. A per-token gate (gated LoRA; Samragh et al., 2025) decides which pathway each token takes, so clean and noisy tokens can be processed in the same forward pass. For a linear layer with frozen weight $W \in \mathbb{R}^{d_{\mathrm{out}} \times d_{\mathrm{in}}}$, LoRA update $\Delta W = BA$ with $B \in \mathbb{R}^{d_{\mathrm{out}} \times r}$, $A \in \mathbb{R}^{r \times d_{\mathrm{in}}}$, and input $h \in \mathbb{R}^{d_{\mathrm{in}}}$,

\[y = W h + g \Delta W h, \qquad g = 1 \text{ for noisy (draft) tokens, } 0 \text{ for clean tokens.}\]

Why LoRA? Speculative decoding is fast only if the draft distribution $q$ closely matches the target distribution $p$. When the drafter shares all of the target’s weights except a low-rank update, the two distributions stay tightly coupled. The adapters are also small: for Qwen3-8B they add 0.35B parameters, compared with 0.40B for the EAGLE-3 drafter (Li et al., 2025) and 1.05B for the DFlash drafter (Chen et al., 2026).

Training the diffusion weights: Diffusion Distillation

Background: discrete diffusion. Discrete diffusion corrupts a clean sequence $\mathbf{x} \in \mathcal{V}^L$ into noise $\mathbf{z} \sim \boldsymbol{\pi}^L$ and learns to undo the corruption. $\pi$ is the noise distribution, i.e., a discrete distribution over $\mathcal{V}$. Masked diffusion uses a special [MASK] token for $\pi$. Uniform-state diffusion uses the uniform distribution over the vocabulary $\mathcal{V}$. We use uniform-state diffusion because it can correct its own mistakes and works well with very few denoising steps (Sahoo et al., 2025).

One step, block by block. We want the diffusion pathway to propose a block of tokens that the AR model would have generated one by one. We take the few-step idea to its extreme and denoise in a single step, from pure noise $\mathbf{z} \sim \mathrm{Unif}(\mathcal{V})^L$ directly to clean tokens. This is a one-step version of discrete consistency distillation (Sahoo et al., 2025), with the frozen AR model as the target. Denoising a long sequence in one step is too hard, so we work in blocks of size $B$: each noisy block is denoised given the clean tokens of all earlier blocks.

Two copies in one forward pass. During training, we run the model once on the concatenation $[\mathbf{x}, \mathbf{z}]$, where:

• $\mathbf{x}$ is the clean sequence of length $L$.
  • It uses the AR pathway (LoRA off).
  • It produces the target distributions.
• $\mathbf{z}$ is a fully noisy sequence of random tokens, also of length $L$.
  • It uses the diffusion pathway (LoRA on).
  • It produces the draft distributions.

Partition the sequence of length $L$ into blocks of size $B$. For a token at index $\ell$ within block $b$,

\[p_{\ell} = p_{\theta_{\mathrm{AR}}}\left(\cdot \mid \mathbf{x}_{< \ell} \right), \qquad q_{\ell} = p_{\theta_{\mathrm{AR}} + \theta_{\Delta}} \left(\cdot \mid \mathbf{x}_{\le (b-1)\cdot B},\ \mathbf{z}_{(b - 1) \cdot B +1 : \ell - 1 }\right).\]

Unlike typical diffusion language models, both pathways retain the next-token parameterization. While drafting a block of tokens, each token attends to clean tokens of earlier blocks and all prior noisy tokens in its own block, which is exactly what it sees at inference as well.

A block-causal attention mask makes this work (Figure 4). Clean tokens never attend to noisy tokens, so the target distributions are not affected by the noise. All blocks are trained in one pass of length $2L$, regardless of the block size. For comparison, training the DFlash drafter needs a context of length $B \cdot L$.

The loss. Only $\theta_{\Delta}$ is trained; the AR weights are frozen. The objective combines the distillation loss (Sahoo et al., 2025) with a total variation (TV) term:

\[\mathcal{L}(\theta_{\Delta}) = \mathbb{E}_{\mathbf{x} \sim \mathcal{D}, \mathbf{z} \sim \mathrm{Unif}(\mathcal{V})^L} \sum_{\ell=1}^{L} \Big[ \alpha \mathrm{KL} \left(q_\ell || p_\ell\right) + \beta \big\lVert q_\ell - p_\ell \big\rVert_1 \Big],\]

where $\mathcal{D}$ is the training data, $\alpha, \beta \ge 0$ weight the two terms, and $\lVert q_\ell - p_\ell \rVert_1 = 2 \mathrm{TV}(q_\ell, p_\ell)$. The TV term targets the acceptance probability directly: the probability that a draft token from $q$ is accepted is exactly $1 - \mathrm{TV}(p, q)$ (Leviathan et al., 2023). Our 7B model uses $\alpha = 0.01$ and $\beta = 1$; our Qwen model uses TV only. We also grow the block size during training, for example, $B = 2 \to 4 \to 8$.

We call this training recipe Diffusion Distillation, derived from Discrete Consistency Distillation (Sahoo et al., 2025), as we distill the AR model into a single-step diffusion over the draft block.

Uno training is cheap. For our 7B model, Uno was trained on around 7B tokens of SFT data, about 0.03% of the AR pretraining data. The phase fits right after SFT and before RL (Figure 1, top), so the adapters can also speed up RL rollouts.

Uno inference with Ψ-Spec samplers

To sample from Uno, we introduce the Ψ-Spec (Ψ-Speculative) sampler. It drafts with Ψ-samplers (Deschenaux et al., 2026), a family of discrete diffusion samplers; with the single denoising step used in all our experiments, this is just sampling from the denoiser. Uno alternates between the following two types of forward passes:

1. Draft. We append $B - 1$ random tokens to the last generated token; this creates a draft block of $B$ tokens. One forward pass with gated LoRA returns $B$ tokens. The first comes from the AR weights alone, so it is an exact sample from the AR model. The other $B - 1$ are drafts $d_j \sim q_j$.
2. Verify. One forward pass with the AR weights alone scores all drafts at once and gives the target distributions $p_j$. We then apply the standard rejection rule of speculative decoding, from left to right. Each draft is accepted with probability

\[\min \left(1,\ \frac{p_j(d_j)}{q_j(d_j)}\right).\]

At the first rejection, we sample a replacement token from the residual distribution and stop:

\[\frac{\max(0,\ p_j - q_j)}{\lVert \max(0,\ p_j - q_j) \rVert_1}.\]

If all drafts are accepted, we sample one extra token from the AR model. This rule guarantees that the output follows exactly the AR model’s distribution (Leviathan et al., 2023; Chen et al., 2023). The guarantee holds however good or bad the drafts are; better drafts only make Uno faster.

How many tokens per forward pass? Each round of draft+verify emits the exact AR token, $n$ accepted drafts (with $0 \le n \le B - 1$), and one corrected or extra token, using two forward passes. The number of tokens per forward pass (TPF) is therefore

\[1 \le \mathrm{TPF} = \frac{n + 2}{2} \le \frac{B + 1}{2}.\]

The worst case is the AR rate: even if every draft is rejected, two passes still give two tokens. Below, we also report $\tau$, the number of tokens per draft-plus-verify step, so $\mathrm{TPF} = \tau / 2$ for Uno.

Tree sampler. The linear sampler, as described above, drafts one candidate continuation per sequence. It is best at large batch sizes, where the GPU is compute-bound. At small batch sizes the GPU is memory-bound, and the tree sampler uses its idle compute: it takes the top $K$ tokens at each draft position, keeps the $V$ most likely prefixes, and verifies the whole tree in one pass (Cai et al., 2024; Li et al., 2025). We use the linear sampler with $B = 4$ at the largest batch, and the tree sampler with $(B, K, V) = (16, 32, 32)$ at batch size 1.

Results

We measure speed with a fixed throughput test: a random 1K-token prompt and 8K output tokens, on one H200 GPU in bfloat16. Each method runs at the average acceptance rate it reaches on real benchmarks. This puts all methods on the same input and output lengths; timing them directly on benchmarks would favor models that write shorter answers. We report per-request throughput at batch size 1, and system throughput, the total tokens per second at the largest batch that fits on the GPU.

Uno initialized from Qwen3-8B

We first train Uno adapters on the open-weight Qwen3-8B. As Qwen’s training data is not public, we train the adapters (rank 128, 0.35B parameters) on a different dataset, OpenThoughts3, for 14.7B tokens. We compare with two lossless speculative decoding methods that provide drafters for Qwen3: EAGLE-3, which drafts with a small AR head, and DFlash, which drafts with a small block-diffusion model. All methods run at temperature 1 with thinking mode on.

Uno is the fastest method across all batch sizes: $1.6\times$ faster than the AR model at the largest batch and $2.5\times$ faster at batch size 1. The tokens per step $\tau$ alone flatter Uno, because Uno drafts with the full network while the baselines draft with small models. Throughput includes this cost, and Uno still wins at every batch size we tested (Figure 6). Uno also needs less memory because the drafter and verifier share one set of weights and one KV cache.

We also compare with lossy parallel decoders built on similar models, such as SDAR, TiDAR, I-DLM, Jacobi Forcing, and FLARE. These methods fine-tune the AR weights, and most of them lose accuracy relative to their parent model, sometimes a lot. For example, SDAR, fine-tuned from Qwen3-8B, scores 10.0% on AIME-25 and 16.6% on LiveCodeBench v6, while Qwen3-8B itself scores 70.7% and 50.9%.

Training from scratch: Uno 7B

Next, we train a diffusion-augmented LLM end to end. The AR weights come from our 7B K2-Horizon model, a dense transformer trained on about 23T tokens with context lengths up to 512K tokens. The diffusion weights are trained after SFT. We compare with three diffusion LLMs: the open DiffusionGemma (26B mixture-of-experts, 4B active parameters; DiffusionGemma Team, 2026), the open Nemotron-Labs-Diffusion (14B; Fu et al., 2026), and the proprietary Mercury 2.

Quality. Since Uno is lossless, its accuracy is the accuracy of its AR model. It beats all three diffusion LLMs on every agentic tool-use, coding, and long-context benchmark we evaluated, wherever baseline results are available. On SWE-bench Verified, Uno scores 70.6% versus 18.7% for DiffusionGemma. On the long-context benchmark AA-LCR, it scores 68.0% versus 36% for Mercury 2 and 19.7% for DiffusionGemma (Figure 1). The one benchmark where it trails is AA-Omniscience, a knowledge-recall test.

Speed. At the largest batch size, Uno reaches 5255 tok/s: $1.5\times$ its own AR model, and more than $4\times$ DiffusionGemma and Mercury 2 (Figure 7). At batch size 1, Uno is $2.2\times$ faster than its AR model. DiffusionGemma and Mercury 2 are faster in this regime (about $4-4.5\times$ AR), but they score lower on most benchmarks, and their reported system throughput is less than a quarter of Uno’s. Both open diffusion baselines are in fact reported to be slower than their own AR base models at large batch sizes (Fu et al., 2026; DiffusionGemma Team, 2026).

The adapters survive RL. We post-trained the AR model with RL after training the Uno adapters. We trained only the AR weights with DAPO to obtain math, code, tool-use, and search experts, and merged the experts into one model. As we kept the diffusion weights frozen in this stage, the drafts might drift away from the verifier. However, with the same adapters, the merged model’s average TPF was 2.10, compared with 2.25 for the SFT model. This shows that the same Uno adapters can be used after RL post-training. In our experiments, using them to generate RL rollouts made RL training up to 40% faster for the math and code experts; tool use and search gained less, since tool calls dominate their runtime.

Limitations and outlook

Uno uses two full-size forward passes per iteration, which caps TPF at $(B+1)/2$. Quadratic sampling (Samragh et al., 2025) could merge drafting and verification into one pass, but it needs efficient kernels to pay off. Like other speculative methods, Uno’s speedup shrinks when the GPU becomes compute-bound: on Qwen3-8B it goes from $2.5\times$ at batch size 1 to $1.6\times$ at batch size 64.

A more open question is inference-time scaling. Ψ-Spec can run more denoising steps per block than it drafts tokens, which spends more compute without making the context longer, and more steps are known to improve sample quality for these samplers. If the diffusion pathway’s samples ever became better than the AR model’s, one could even turn verification off. We leave this to future work.

References

• Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., and Dao, T. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. ICML 2024.
• Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arXiv:2302.01318.
• Chen, J., Liang, Y., and Liu, Z. (2026). DFlash: Block Diffusion for Flash Speculative Decoding. arXiv:2602.06036.
• Deschenaux, J., Gulcehre, C., and Sahoo, S. S. (2026). The Diffusion Duality, Chapter II: Ψ-Samplers and Efficient Curriculum. ICLR 2026.
• DiffusionGemma Team (2026). DiffusionGemma Technical Report. arXiv:2608.00146.
• Fu, Y., Whalen, L., Garg, A., et al. (2026). Nemotron-Labs-Diffusion: A Tri-Mode Language Model Unifying Autoregressive, Diffusion, and Self-Speculation Decoding. arXiv:2607.05722.
• Gong, S., Agarwal, S., Zhang, Y., Ye, J., Zheng, L., Li, M., An, C., Zhao, P., Bi, W., Han, J., Peng, H., and Kong, L. (2025). Scaling Diffusion Language Models via Adaptation from Autoregressive Models. ICLR 2025.
• Kim, M., Hooper, C., Tomar, A., Xu, C., Farajtabar, M., Mahoney, M. W., Keutzer, K., and Gholami, A. (2025). Beyond Next-Token Prediction: A Performance Characterization of Diffusion versus Autoregressive Language Models. arXiv:2510.04146.
• Leviathan, Y., Kalman, M., and Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. ICML 2023.
• Li, Y., Wei, F., Zhang, C., and Zhang, H. (2025). EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test. NeurIPS 2025.
• Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.-R., and Li, C. (2025). Large Language Diffusion Models. NeurIPS 2025.
• Sahoo, S. S., Deschenaux, J., Gokaslan, A., Wang, G., Chiu, J. T., and Kuleshov, V. (2025). The Diffusion Duality. ICML 2025.
• Samragh, M., Kundu, A., Harrison, D., Nishu, K., Naik, D., Cho, M., and Farajtabar, M. (2025). Your LLM Knows the Future: Uncovering Its Multi-Token Prediction Potential. arXiv:2507.11851.
• Ye, J., Xie, Z., Zheng, L., Gao, J., Wu, Z., Jiang, X., Li, Z., and Kong, L. (2025). Dream 7B: Diffusion Large Language Models. arXiv:2508.15487.