TLDR: dInfer is a novel, efficient, and extensible inference framework for Diffusion-based Large Language Models (dLLMs). It addresses key bottlenecks like high computational cost and limited parallel decoding efficiency through a modular design, integrating innovative algorithms for diffusion iteration management, decoding strategies (hierarchical, credit), and KV-cache management (vicinity refresh). Coupled with system-level optimizations like tensor/expert parallelism, PyTorch compilation, and loop unrolling, dInfer achieves substantial speedups (e.g., 10x over Fast-dLLM, 2-3x over vLLM with QWen2.5-3B) while maintaining output quality, making dLLMs more practical for deployment.
Diffusion-based Large Language Models, or dLLMs, are emerging as a powerful new approach in the world of artificial intelligence, offering an alternative to the more traditional autoregressive (AR) models. Unlike AR models that generate text word by word, dLLMs refine entire sequences in parallel through an iterative denoising process. This inherent parallelism holds immense promise for faster processing and better utilization of powerful GPU hardware.
Despite their potential, dLLMs have faced significant hurdles in widespread adoption. Key challenges include their computational intensity due to iterative denoising, difficulties in scaling parallel decoding without compromising output quality, and a lack of standardized, efficient inference frameworks. These bottlenecks have made it hard to consistently benchmark and deploy dLLMs effectively.
Introducing dInfer: A New Framework for dLLMs
To tackle these challenges, researchers have introduced dInfer, an efficient and extensible inference framework specifically designed for dLLMs. dInfer breaks down the complex inference pipeline into four manageable, modular components: the model itself, a diffusion iteration manager, a decoding strategy, and a KV-cache manager. This modularity allows for flexible combinations of algorithms, enabling users to create customized inference pipelines that maximize the benefits of parallel decoding while significantly boosting computational efficiency.
Algorithmic Innovations for Enhanced Performance
dInfer integrates several novel algorithms to achieve its efficiency gains. The diffusion iteration manager, which controls the iterative denoising, includes an ‘iteration smoothing’ algorithm. This innovation reuses information from previous steps, enriching contextual cues and boosting token confidence, especially when KV-cache is used.
For parallel decoding, dInfer offers three strategies: ‘threshold decoding’ (committing tokens above a certain confidence), ‘hierarchical decoding’ (recursively partitioning masked sections to reduce local dependencies and improve efficiency), and ‘credit decoding’ (accumulating historical confidence scores to prioritize consistently stable predictions, thereby improving reliability). These strategies allow dInfer to decode more efficiently without needing to retrain the underlying dLLM.
A major challenge for dLLMs is the incompatibility with traditional KV-cache reuse, which is common in AR models. dLLMs’ token representations evolve, making static cache reuse ineffective. dInfer addresses this with ‘vicinity KV-cache refresh,’ which intelligently updates a small window of tokens adjacent to the current decoding block, balancing cost and performance by exploiting semantic locality.
System-Level Optimizations for Speed
Beyond algorithmic improvements, dInfer incorporates several system-level optimizations. It supports both tensor parallelism (TP) and expert parallelism (EP) to maximize GPU utilization, even with small batch sizes. This combination can improve inference efficiency by over 100%.
To further optimize single-sequence inference, dInfer leverages PyTorch’s just-in-time (JIT) compiler and NVIDIA CUDA Graphs. This eliminates PyTorch execution overhead and fuses CUDA kernels, leading to over 200% improvement in inference efficiency when TP and EP are enabled. Additionally, a ‘loop unrolling’ mechanism addresses CUDA stream bubbles, ensuring continuous kernel launches and boosting iteration efficiency by 5-10%. An early termination mechanism also prevents unnecessary computations once an end-of-sequence token is generated.
Also Read:
- DiTSinger: Advancing Singing Voice Synthesis with Scalable Data and Implicit Alignment
- Dynamic Temperature Control Enhances LLM Reasoning in Reinforcement Learning
Impressive Performance Benchmarks
Evaluations show dInfer’s significant impact. On the HumanEval benchmark, dInfer achieves over 1,100 tokens per second (TPS) at batch size 1, and averages more than 800 TPS across six benchmarks on a single node with 8× H800 GPUs. Compared to Fast-dLLM, a prior system, dInfer delivers more than a 10× speedup while maintaining accuracy. It also provides a 2–3× speedup over highly optimized AR models like QWen2.5-3B running on the vLLM inference engine, with comparable output quality.
The framework also supports enhanced dLLM variants like LLaDA-MoE-TD, which uses a technique called Trajectory Distillation to fine-tune models for even greater parallel decoding efficiency. This distilled model, combined with dInfer’s optimizations, achieves an average TPS of 847.22, more than 3x faster than Qwen2.5-3B in vLLM.
dInfer represents a significant step forward for dLLMs, providing a robust and efficient framework that not only accelerates inference but also offers a standardized platform for future research and development in this rapidly evolving field. The implementation of dInfer is open-sourced and available for exploration. You can find the research paper here: dInfer: An Efficient Inference Framework for Diffusion Language Models.
