dlmrt runs one request at a time and makes that request fast. Every forward denoises a whole block of tokens, and the denoising loop runs on the GPU. It serves QwenDLM, LLaDA2.1-mini and DiffusionGemma checkpoints on one GPU or tensor parallel across an NVLink node, on NVIDIA Hopper (H100) and, for LLaDA2.1-mini and DiffusionGemma, Blackwell (B200, B300; NVFP4 also on the RTX 5090). It is the inference engine of On-Policy Residual Learning for Diffusion Language Models (ORL).

Highlights

• Denoising on the device. Each block is one CUDA graph, begin → WHILE(not done){round} → commit, with a conditional node that decides every round on the GPU. The host launches a block and reads back its tokens.
• The papers' numerics. Kernels follow each reference implementation's BF16 roundings. For QwenDLM and LLaDA2, the SGLang-numerics mode runs the paper runtime's forward pass and decisions: 4,816 of the 4,820 QwenDLM answers and all 164 LLaDA2.1-mini ORL Speed HumanEval answers are identical to the paper runtime's.
• Tensor parallel without barriers. A fused NVLink reduction (push reduce-scatter, rank-ordered FP32 sum, multicast all-gather, in-band arrival detection) returns identical values on every rank, so all ranks run the same device loop with no host coordination.
• Overlapped launches. Every Triton kernel and the patched DeepGEMM GEMMs use programmatic dependent launch: each kernel's launch overlaps its predecessor's tail.
• Low precision. FP8 block-scaled projections for QwenDLM; NVFP4 experts on Blackwell FP4 tensor cores for LLaDA2 and DiffusionGemma (FlashInfer's TRT-LLM-gen MoE; on RTX 50-series GPUs dlmrt's Triton MoE with the TRT-LLM-gen MoE's roundings for LLaDA2, bitwise the B200's in all but about 0.2% of rows, and FlashInfer's CUTLASS MoE for DiffusionGemma), with an FP8 KV cache.

Supported models

Each family's decoders, precisions and options: docs/models.md.

Performance

ORL checkpoints and decoders: the latency of one denoise round at an empty context, and request TPS on full HumanEval.

Each model runs the paper's decoder (LLaDA2: Quality mode; DiffusionGemma: 2048 new tokens). The TP=1 TPS are the paper's: on H100, QwenDLM and LLaDA2 run the SGLang-numerics mode and DiffusionGemma averages seeds 1235, 1239 and 1241; on one B200, DiffusionGemma averages seeds 1234–1236. The round latencies and TP=8 run dlmrt's default numerics (DiffusionGemma at seed 1234), which reach 1658 (QwenDLM), 2060 (LLaDA2) and 1388 (DiffusionGemma) request TPS at TP=1 on H100. Methodology and every measurement: docs/performance.md.

Installation

dlmrt runs on Linux with Python 3.10+ and compiles its CUDA extensions on first use. It needs PyTorch 2.11, Triton 3.6 and a CUDA toolkit (12.9+ for Blackwell, 13.0+ for B300).

The SGLang image lmsysorg/sglang:v0.5.13 ships every dependency; install dlmrt into its system Python:

git clone https://github.com/Introspective-Diffusion/DLMRT && cd DLMRT
pip install --no-deps --break-system-packages -e .

In another environment, pip installs the dependencies:

pip install -e .                    # torch, triton, DeepGEMM (sgl-deep-gemm), FA3 (sglang-kernel), transformers
pip install -e ".[blackwell]"       # + FlashInfer: NVFP4 experts on Blackwell
pip install -e ".[sglang-numerics]" # + FlashInfer 0.6.12: the QwenDLM / LLaDA2 SGLang-numerics modes
pip install -e ".[bench]"           # + pyarrow: dlmrt bench humaneval, benchmarks/humaneval_score.py

Quickstart

Python

from dlmrt import LLM

llm = LLM("google/diffusiongemma-26B-A4B-it")  # a local directory or a Hugging Face repo id
out = llm.generate("Write a Python function that checks whether a number is prime.", max_new_tokens=2048)
print(out.text)
print(f"{out.request_tps:.0f} tok/s, {sum(out.denoise_steps)} denoise rounds")

For tensor parallelism, run the same program under torchrun --nproc-per-node 8: every rank builds the same LLM and makes the same calls (examples/offline_inference.py).

Server

Streaming NDJSON over HTTP, one request at a time (docs/server.md):

torchrun --nproc-per-node 8 -m dlmrt serve --model CKPT --port 18080 [--role orl] [--precision nvfp4]
curl -N localhost:18080/generate -d '{"messages": [{"role": "user", "content": "Hi!"}], "max_new_tokens": 512}'

Benchmarks

python -m dlmrt bench step --model CKPT    # round latency at fixed contexts
# HumanEval generation (QwenDLM, DiffusionGemma; LLaDA2: docs/benchmarks.md)
torchrun --nproc-per-node 8 -m dlmrt bench humaneval --model CKPT --inputs openai_humaneval.parquet \
    --out results/he --full --repeats 1 --max-new-tokens 2048 [--role orl]
python benchmarks/humaneval_score.py openai_humaneval.parquet results/he   # runs model-written code: use a sandbox

Reproducing the papers

dlmrt bench suite runs the papers' generation protocols over their problem files. The commands for every row of the results table, the settings and the scoring are in docs/benchmarks.md.

Documentation

Repository layout

dlmrt/
  llm.py           LLM: the Python API (also used by the server)
  models/          sdar/, llada2/, diffusion_gemma/: config, model and engine per family, kernels where it has its own
  layers/          Linear (DeepGEMM BF16 / FP8, cuBLAS, a Triton GEMM in cuBLAS's B200 order)
  quantization/    NVFP4 (ModelOpt exports): dequantization, expert layouts for the TRT-LLM-gen, CUTLASS and Triton MoEs
  distributed/     the tensor / expert parallel group and the fused reduction
  runtime/         shared engine pieces (Generation, check_prompt)
  ops/             triton/ shared kernels, csrc/ CUDA sources, extensions.py native loader,
                   fa3_tma.py FlashInfer's FA3 with the paged kernel's numerics (SGLang-numerics modes)
  entrypoints/     `dlmrt bench`, `dlmrt serve`
  envs.py          every DLMRT_* environment variable
benchmarks/        scoring, TPF / TPS of suite records, token-identity comparison, kernel sweeps
examples/          offline inference
tests/             kernel tests and checkpoint-gated parity tests
tools/             check_imports.py (the static import check CI runs), architecture_svg.py (the block-graph figure)

Citation

@misc{yu2026orl,
  title         = {On-Policy Residual Learning for Diffusion Language Models},
  author        = {Yu, Yifan and Jian, Yuqing and Li, Jisen and Hu, Yuezhou and Zhang, Tianyi and Wu, Xiaoxia and
                   Zhu, Rui-Jie and Wang, Junxiong and Zhou, Zhongzhu and Lai, Fan and Zou, James and Xu, Chenfeng},
  year          = {2026},
  eprint        = {2610.20687},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2610.20687}
}

License

Apache License 2.0 (LICENSE). Vendored third-party code keeps its license: ThunderKittens and a modified DeepGEMM kernel (MIT), kernels adapted from SGLang (Apache 2.0); see THIRD_PARTY_NOTICES.md.