WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference
Abstract
Autoregressive (AR) generation is the standard decoding paradigm for Large Language Models (LLMs), but its token-by-token nature limits parallelism at inference time. Diffusion Language Models (DLLMs) offer parallel decoding by recovering multiple masked tokens per step; however, in practice they often fail to translate this parallelism into speed gains over optimized AR engines (e.g., vLLM). A key reason is that many DLLMs rely on bidirectional attention, which breaks standard prefix KV caching. We propose WeDLM, a diffusion decoding framework built entirely on standard causal attention to make parallel generation prefix-cache friendly. The core idea is to let each masked position condition on all observed tokens while keeping a causal mask, achieved by Topological Reordering that moves observed tokens to the physical prefix while preserving their logical positions. Building on this, we introduce a streaming decoding procedure that continuously commits confident tokens into a growing left-to-right prefix, avoiding the stop-and-wait behavior common in block diffusion methods. Experiments show that WeDLM preserves the quality of strong AR backbones while delivering substantial speedups, approaching 3× on challenging reasoning benchmarks and up to 10× in low-entropy generation regimes; critically, our comparisons are against AR baselines served by vLLM under matched deployment settings.
Lay Summary
Modern AI assistants like ChatGPT generate text one word at a time. Because each word must wait for the previous one to be produced, this sequential process limits how fast the system can respond, even on powerful parallel hardware. Researchers have explored an alternative called "diffusion" that writes several words at once. In theory this should be much faster. In practice, however, the most popular diffusion methods are not actually faster than standard left-to-right generators — their internal design prevents them from reusing earlier computations between steps, so the model ends up repeating work it has already done. Our work, WeDLM, fixes this. We redesigned how the model reads its context so it can still write several words in parallel each step, but in a way that lets it save and reuse earlier work instead of recomputing it. The result is an AI text generator that is up to three times faster than a strong standard baseline on math and reasoning tasks, and up to ten times faster on highly predictable outputs, while producing answers of equal or higher quality — bringing parallel text generation closer to being practical for real deployment.