Unlocking Lossless Speedups in LLMs
via Discrete Diffusion

Subham Sekhar Sahoo, Lingjie Chen, Khiem Pham, Jonathan Geuter, Junlin Chen,
Chaitanya Dwivedi, Varad Pimpalkhute, Yash Akhauri, Alexander Moreno, Mikhail Yurochkin,
Zhenting Wang, Mostafa Elhoushi, Nolan Dey, Shane Bergsma, Joel Hestness,
Hongyi Wang, John Thickstun, Eric Xing, Zhengzhong Liu

We introduce Uno, a diffusion-augmented LLM that features two sets of weights:

• AR weights: Trained using next-token prediction loss to define an AR distribution.
• Diffusion weights: Trained to generate multiple tokens in parallel from the AR distribution.

To sample from Uno, we propose $\Psi$-Spec sampler, which enables provably lossless multi-token prediction from the AR distribution.

• Unlike speculative decoding methods, Uno does not require a separately trained draft model and achieves higher throughput across all batch sizes.
• Unlike self-speculative deciding approaches, Uno is lossless.

In this repo, we release the code for:

• Inference
  • A Nano-vLLM-based inference engine for $\Psi$-Spec sampler.
    • Linear Sampler for high system throughput.
    • Tree sampler for high per-request throughput.
  • Ready-to-run recipes for Uno Qwen3 8B, Uno 7B, and Uno 0.9B.
• Training
  • A conditional-LoRA diffusion training pipeline for Uno Qwen3 8B.
  • OpenThoughts data preparation and progressive block-size curricula.
  • Multi-node Slurm launchers with checkpoint resume support.
• Evaluation
  • One shared generation and scoring pipeline across model families.
  • Benchmark-specific data loaders and graders.
  • Accuracy, tokens per forward (TPF), and tokens per second (TPS) metrics.
  • Single-benchmark, full-suite, and persistent pull-based launchers.

Table of Contents

• Code Organization
• Getting Started
  • Installation
  • Checkpoints
• Reproducing Experiments
  • Training
  • Inference
  • Evaluation
• Acknowledgements
• Citation

Code Organization

The repository is organized around the workflows users run:

• nano_vllm_uno/ : Inference engine optimized for Uno, with model execution, KV-cache management, and lossless linear and tree sampling.
• generation.py : Shared model loading, conditional-LoRA attachment, prompt formatting, generation, and TPF/TPS accounting.
• inference.py : Free-form prompt inference with linear or tree sampling.
• evaluation/ : Benchmark protocols, dataset loaders, generation, parsers, and benchmark-specific graders.
• training/ : Data preparation, conditional-LoRA training, diffusion objectives, curricula, and checkpointing.
• examples/ : Model-specific training, inference, and evaluation recipes that call the shared workflows.

Getting Started

Installation

Create a Python 3.10 environment and install Uno with its evaluation and training dependencies:

conda create -n nano-vllm-uno python=3.10 pip -y
conda activate nano-vllm-uno

python -m pip install --upgrade pip
python -m pip install torch==2.11.0 \
  --index-url https://download.pytorch.org/whl/cu128
python -m pip install \
  'https://github.com/lesj0610/flash-attention/releases/download/v2.8.3-cu12-torch2.11/flash_attn-2.8.3%2Bcu12torch2.11cxx11abiTRUE-cp310-cp310-linux_x86_64.whl'
python -m pip install -e '.[eval,train]'

This installs FlashAttention-2 (FA2), which is sufficient for linear decoding. Tree verification additionally requires FlashAttention-3 (FA3):

python -m pip install ninja==1.13.0
git clone --depth 1 --branch v2.8.3 \
  https://github.com/Dao-AILab/flash-attention.git
cd flash-attention/hopper
MAX_JOBS=16 python -m pip install --no-build-isolation .

Checkpoints

Model	Base model	Uno weights
Uno 7B	IFM/K2-Horizon-7B	IFM/K2-Horizon-7B-Uno
Uno 0.9B	IFM/K2-Horizon-0.9B	IFM/K2-Horizon-0.9B-Uno
Uno Qwen3 8B	s-sahoo/uno-qwen3-8B	s-sahoo/uno-qwen3-8B/adapter

Public checkpoints can be downloaded without a Hugging Face token. A token is still required for gated datasets such as GPQA and for any private or gated model repository.

Reproducing Experiments

All public model recipes live under examples/. Each recipe supplies the model-specific defaults and delegates to the same shared training, inference, or evaluation workflow.

Training

To train Uno Qwen3 8B, prepare the training data and run the shared training entry point as follows.

1. Prepare the Training Data

Prepare the pinned OpenThoughts3-1.2M corpus once in shared storage:

python -m training.prepare_openthoughts \
  --output /path/to/openthoughts-uno-4095 \
  --cache-dir /path/to/hf-cache \
  --num-proc 32

2. Train Uno Qwen3 8B

The following is a single-GPU command. Gradient accumulation preserves the released global batch size of 128:

python -m training.train \
  --dataset-path /path/to/openthoughts-uno-4095 \
  --output-dir /path/to/uno-training \
  --curriculum training/configs/uno_3epoch_curriculum.yaml \
  --deepspeed training/configs/deepspeed_zero2.json \
  --per-device-batch-size 8 \
  --gradient-accumulation-steps 16 \
  --learning-rate 1e-5 \
  --warmup-steps 562 \
  --lora-target all \
  --lora-rank 128 \
  --lora-alpha 2048 \
  --ce-alpha 0 \
  --kl-beta 0 \
  --tv-gamma 1

See training/train.py and training/configs/uno_3epoch_curriculum.yaml for the complete training configuration. Resume an interrupted run with --resume-from-checkpoint /path/to/checkpoint and the same --output-dir.

Inference

Use the model recipe that matches the checkpoint. All three recipes call the same shared inference.py workflow:

# Uno Qwen3 8B
bash examples/uno_qwen3_8B/run_inference.sh \
  --prompt "Solve 2 + 2 and explain your reasoning."

# Uno 8B
bash examples/uno_8B/run_inference.sh \
  --prompt "Solve 2 + 2 and explain your reasoning."

# Uno 1B
bash examples/uno_1B/run_inference.sh \
  --prompt "Solve 2 + 2 and explain your reasoning."

The example launchers use linear sampling by default. To enable tree sampling, install FA3 and pass the tree parameters to the same entry point:

ATTENTION_BACKEND=fa3 \
  bash examples/uno_qwen3_8B/run_inference.sh \
    --prompt "Solve 2 + 2 and explain your reasoning." \
    --diffusion-block-size 16 \
    --tree-candidate-top-k 32 \
    --tree-verify-size 60

The command prints generated text together with output-token count, elapsed time, TPS, decoder statistics, and TPF. Common controls include --temperature, --top-p, --top-k, --max-tokens, --diffusion-block-size, and --max-num-seqs.

Evaluation

The evaluation workflow uses canonical benchmark settings from evaluation/benchmarks.py. Generation and grading are separate internally, while each model recipe runs both stages through one command. Evaluation results are provided in each model-specific subdirectory under examples/.

1. Prepare Evaluation Data

Missing public datasets are prepared automatically. They can also be prepared in advance:

export UNO_EVAL_DATA_DIR=/path/to/uno-eval-data

python -m evaluation.prepare_data \
  --output-dir "${UNO_EVAL_DATA_DIR}"

Keep UNO_EVAL_DATA_DIR set when running run_eval.sh. The evaluation runner will reuse the prepared JSONL files from this directory and will prepare any missing benchmark there automatically.

2. Evaluate One Benchmark

Set the shared output and parallelism options, then choose the model recipe:

export RESULTS_ROOT=/path/to/results
export DATA_PARALLEL_SIZE=1

# Uno Qwen3 8B
bash examples/uno_qwen3_8B/run_eval.sh gsm8k

# Uno 8B
bash examples/uno_8B/run_eval.sh gsm8k

# Uno 1B
bash examples/uno_1B/run_eval.sh gsm8k

Supported benchmarks:

• Math: aime24, aime25, aime26, gsm8k, and math500.
• Knowledge: arc_challenge, gpqa_diamond, hle, and omniscience.
• Code: humaneval and mbpp.
• Instruction and long-context reasoning: ifeval and lcr.

3. Evaluate the Full Suite

Run the complete suite sequentially:

MODEL_EXAMPLE=uno_qwen3_8B \
RESULTS_ROOT=/path/to/results \
DATA_PARALLEL_SIZE=1 \
  bash evaluation/run_suite.sh

Set MODEL_EXAMPLE to uno_qwen3_8B, uno_7B, or uno_0.9B.

Acknowledgements

Uno's runtime builds on Nano-vLLM, and its training utilities are adapted from LlamaFactory. See NOTICE for attribution and license notices.

Citation

@misc{sahoo2026unlockinglosslessspeedupsllms,
      title={Unlocking Lossless Speedups in LLMs via Discrete Diffusion},
      author={Subham Sekhar Sahoo and Lingjie Chen and Khiem Pham and Jonathan Geuter and Chaitanya Dwivedi and Varad Pimpalkhute and Yash Akhauri and Alexander Moreno and Mikhail Yurochkin and Zhenting Wang and Mostafa Elhoushi and Nolan Dey and Shane Bergsma and Joel Hestness and John Thickstun and Eric Xing and Zhengzhong Liu},
      year={2026},
      eprint={2609.04010},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2609.04010},
}