Example generation on AIME24. RCD increases parallelism by 4x while maintaining the baseline's peak accuracy.

Introduction

This repository contains the code to replicate our study in "Residual Context Diffusion Language Models". In this study, we point out that diffusion Large Language Models (dLLMs) enable parallel decoding but often trail autoregressive models in accuracy. A key culprit is the inference-time remasking strategy that commits only high-confidence tokens and discards the rest, wasting intermediate computation.

RCD introduces a residual denoising mechanism that turns discarded token distributions into contextual residuals and injects them into the next denoising step. With a two-stage training pipeline, RCD avoids backprop-through-time memory costs while preserving the benefits of residual feedback.

News

• [2026/06] RCD inference is now reference-free: cold-start initialization removes the need for a draft/reference model during generation.
• [2026/06] JetEngine for RCD released: a dedicated CUDA/Triton inference engine for block-diffusion LLMs, optimized for RCD's residual denoising with continuous batching and fused kernels. Supports both RCD and SeqD. See jetengine/.
• [2026/05] RCD is accepted at ICML 2026.
• [2026/02] Project page, arXiv and models are published.

Performance

TL;DR: RCD consistently improves diffusion reasoning accuracy over Sequential Denoising (SeqD) across both SDAR and LLaDA, with the biggest gains on harder competition-style benchmarks (AIME24/25) and MinervaMath.

SDAR

• Models: SDAR 4B / 8B, block size b=32 / 64 (KV cache reuse)
• Eval: SeqD/RCD use 16,384 sequence length; Chat uses 512 tokens (and 1,024 for AIME); confidence threshold = 0.85

LLaDA

• Eval: sequence length 512, single-token-per-step decoding

Model Zoo

All checkpoints are available on HuggingFace.

SeqD Models (Sequential Denoising — Baseline)

RCD Models (Residual Context Diffusion — Ours)

Since June 2026, RCD inference uses cold-start initialization and no longer requires a separate draft/reference model at generation time.

Quick Start

The standalone generation scripts require only transformers:

pip install transformers==4.52.3

# Sequential Denoising (baseline)
CUDA_VISIBLE_DEVICES=0 python seqd-sdar/generate.py \
  --model_dir yuezhouhu/SeqD-SDAR-4B-b64-Thinking \
  --trust_remote_code \
  --block_length 64 \
  --denoising_steps 64 \
  --temperature 0 \
  --dtype bfloat16 \
  --confidence_threshold 0.85

# Residual Context Diffusion (ours)
CUDA_VISIBLE_DEVICES=0 python rcd-sdar/generate.py \
  --model_dir yuezhouhu/RCD-SDAR-4B-b64-Thinking \
  --trust_remote_code \
  --block_length 64 \
  --denoising_steps 64 \
  --temperature 0 \
  --dtype bfloat16 \
  --confidence_threshold 0.85

Repository Layout

residual-context-diffusion/
├── rcd-sdar/                  # RCD on SDAR backbone
│   ├── generate.py            #   Standalone generation
│   ├── lighteval/             #   Evaluation harness
│   ├── recipes/               #   Training recipes
│   ├── eval_simple.sh         #   Eval: GSM8K + MATH500
│   └── eval_aime.sh           #   Eval: AIME24 + AIME25
│
├── seqd-sdar/                 # SeqD on SDAR backbone (baseline)
│   ├── generate.py
│   ├── lighteval/
│   ├── recipes/
│   ├── eval_simple.sh
│   └── eval_aime.sh
│
├── rcd-llada/                 # RCD on LLaDA backbone
│   ├── examples/              #   Eval and training examples
│   └── dllm/                  #   dLLM model code
│
├── seqd-llada/                # SeqD on LLaDA backbone (baseline)
│   ├── examples/
│   └── dllm/
│
├── jetengine/                 # High-performance inference engine
│   ├── rcd/                   #   RCD engine (with residual guidance)
│   ├── seqd/                  #   SeqD engine (baseline)
│   └── README.md
│
└── assets/                    # Figures and diagrams

Installation

Each sub-project is self-contained and has its own environment. Follow the per-directory README for setup:

• RCD SDAR: rcd-sdar/README.md
• SeqD SDAR: seqd-sdar/README.md
• RCD LLaDA: rcd-llada/README.md
• SeqD LLaDA: seqd-llada/README.md

Evaluation

• SDAR: rcd-sdar/eval_simple.sh / rcd-sdar/eval_aime.sh (RCD) and corresponding scripts under seqd-sdar/
• LLaDA: rcd-llada/examples/llada/eval_openmathinstruct.sh and seqd-llada/examples/llada/eval_openmathinstruct.sh

High-Performance Inference (JetEngine)

For batched, high-throughput inference or RL training loops, see jetengine/README.md. JetEngine provides:

• Paged KV-cache with block-level allocation
• Fused Triton attention and MoE kernels
• Native block-diffusion scheduling
• Residual hidden-state support (RCD engine only)

Training

Training recipes are in each method's directory:

• SDAR: rcd-sdar/recipes/ (RCD) and seqd-sdar/recipes/ (SeqD)
• LLaDA: rcd-llada/examples/llada/ and seqd-llada/examples/llada/

Citation

@misc{hu2026residualcontextdiffusionlanguage,
      title={Residual Context Diffusion Language Models}, 
      author={Yuezhou Hu and Harman Singh and Monishwaran Maheswaran and Haocheng Xi and Coleman Hooper and Jintao Zhang and Aditya Tomar and Michael W. Mahoney and Sewon Min and Mehrdad Farajtabar and Kurt Keutzer and Amir Gholami and Chenfeng Xu},
      year={2026},
      eprint={2601.22954},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2601.22954}, 
}

Footnotes

1. We observed potential data contamination in the original Chat models on GSM8K, which may inflate the Chat baseline. ↩

2. Chat variants are instruction-following models; SeqD/RCD are further adapted for mathematical reasoning. ↩ ↩² ↩³ ↩⁴

3. LLaDA-Base results are taken from the original LLaDA paper. ↩