ORL 让扩散语言模型单张 RTX 5090 达 2,788 tokens/s,代码模型全部开源
[Research] 2,788 tokens/s on a single RTX 5090 with a diffusion LLM — Introducing ORL (fully open source) [R]
研究团队发布 On-Policy Residual Learning(ORL)方法,让扩散语言模型在不牺牲生成质量的前提下显著加速,配合自研推理引擎 dlmrt 在单张 RTX 5090 上实现 2,788 tokens/s,8× B200 上 DiffusionGemma 26B MoE 达 7,708 tokens/s。
| Hey r/MachineLearning! We're excited to share ORL (On-Policy Residual Learning), our new research on making diffusion language models (dLLMs) significantly faster without sacrificing generation quality. With ORL and our custom inference engine dlmrt, we've achieved:
All numbers are for single-request decoding (batch size 1), not aggregate throughput across concurrent requests. The GPU figures above are peak results on specific benchmarks, not averages across arbitrary prompts. We've open-sourced the models, training code, evaluation pipeline, and inference engine. What's the idea behind ORL?Unlike autoregressive LLMs, which generate one token at a time, diffusion LLMs can refine and finalize multiple tokens in parallel. But getting them to converge in fewer forward passes is challenging. Our approach has two key ingredients: 1. On-policy residual learning Instead of predicting everything from scratch at each denoising step, the model carries forward its previous predictions as residual context and progressively refines them. The catch: these residuals are generated by the model itself, and their distribution changes during training. Training on stale, offline residuals creates a mismatch with inference. ORL addresses this by training directly on fresh trajectories generated by the current model, replaying the exact residual states encountered during inference. 2. Learning to reduce uncertainty faster Simply supervising the final output doesn't necessarily teach the model to converge faster. We introduce Adjacent Entropy Contraction, which encourages each denoising step to reduce uncertainty relative to the previous step. We also gate the loss to avoid unnecessarily pushing already-confident predictions. The result: the model learns to finalize more tokens per forward pass while maintaining or improving accuracy. Some benchmark resultsWe tested ORL on three model families:
For a more direct comparison with autoregressive decoding, on a single H100, our QwenDLM-4B + ORL reaches 1,298 tokens/s, versus 746 tokens/s for DFlash and 405 tokens/s for EAGLE3 under relaxed verification at comparable measured accuracy. These are end-to-end system comparisons: ORL uses dlmrt, while the speculative decoders use vLLM. We also built a custom inference engine: dlmrtWe developed dlmrt, optimized specifically for low-latency, single-request diffusion decoding. Instead of launching and synchronizing every denoising step from the CPU, dlmrt captures the entire denoising loop in a CUDA graph and executes it on the GPU. It also supports NVFP4 inference on Blackwell GPUs, including the consumer RTX 5090. With identical models and decoding settings, dlmrt runs 1.2–1.6× faster than SGLang/vLLM in our single-H100 comparisons. On the training side, we also co-designed the system around efficient residual-state capture, specialized attention kernels, and compute-balanced packing, accelerating training steps by up to 5.2×. Everything is open source
If you have an RTX 5090, you can try the LLaDA-2.1-mini ORL NVFP4 checkpoint with dlmrt. We'd love feedback from the community, especially anyone interested in reproducing the results, testing different workloads, or contributing to the inference engine. Curious to hear what you think about diffusion LLMs for local inference. Is single-request latency the main thing you'd want optimized, or are there other bottlenecks you'd prioritize? submitted by /u/arty_photography[link] [留言] |
来源:r/MachineLearning · reddit.com