WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference

TL;DR

WeDLM tackles the bottleneck of autoregressive decoding by reconciling diffusion-style parallel mask recovery with standard causal attention. The core ideas are Topological Reordering to expose all observed context under a causal mask and Dual-Stream Masking to align training with inference, followed by Streaming Parallel Decoding that commits a growing left-to-right prefix while maintaining a fixed parallel workload. Empirically, WeDLM preserves or improves the capabilities of strong AR checkpoints and achieves substantial speedups, including up to approximately 3x end-to-end speedups on complex reasoning tasks and up to 10x in low-entropy generation, under matched deployment with optimized AR engines. The work reframes efficiency around prefix cacheability, showing that maximum throughput comes from rapid, cache-friendly prefix growth rather than merely increasing tokens-per-forward, with practical implications for deploying diffusion-based decoders on industrial inference systems.

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 deployment 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 and forces repeated contextualization, undermining efficiency. 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 currently observed tokens while keeping a strict causal mask, achieved by Topological Reordering that moves observed tokens to the physical prefix while preserving their logical positions. Building on this property, we introduce a streaming decoding procedure that continuously commits confident tokens into a growing left-to-right prefix and maintains a fixed parallel workload, 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 3x on challenging reasoning benchmarks and up to 10x in low-entropy generation regimes; critically, our comparisons are against AR baselines served by vLLM under matched deployment settings, demonstrating that diffusion-style decoding can outperform an optimized AR engine in practice.

Figures (8)

• Figure 1: Performance and capability overview of WeDLM-8B. (a) Speed vs. Accuracy:WeDLM-8B achieves a $\sim$3$\times$ speedup over the vLLM-optimized AR baseline (Qwen3-8B) on GSM8K, while also significantly outperforming prior diffusion models in both inference speed (tps) and accuracy. Dream and LLaDA use the dInfer inference engine ma2025dinfer; SDAR uses JetEngine. (b) Holistic Evaluation:WeDLM-8B-Instruct matches or surpasses the strong capabilities of the Qwen3-8B-Instruct baseline, showing improvements across several mathematical, coding, and general knowledge benchmarks.
• Figure 2: Overview of the WeDLM training framework.Left:Topological Reordering physically shifts observed tokens to the prefix while preserving logical positions. This grants masked tokens access to the full observed context under standard causal masking. Right:Dual-Stream Masking concatenates a clean Memory Stream with a masked Prediction Stream. The block-wise attention mask ensures that the Prediction Stream conditions on clean memory history rather than noisy preceding predictions, aligning training dynamics with inference.
• Figure 3: Block Decoding vs. WeDLM Streaming Parallel Decoding. Block decoding suffers from stop-and-wait: bidirectional dependence within a block prevents committing any token until the entire block is finalized. In contrast, WeDLM uses standard causal attention with a dynamic sliding window: resolved tokens (e.g., A, B) are immediately cache-ready and committed, while new mask tokens (e.g., C, E) are appended for parallel prediction.
• Figure 4: Ablation studies on inference hyperparameters. (a) Effect of entropy threshold $\tau$ on MATH accuracy and generation speed, revealing a quality-speed trade-off with optimal range $\tau \in [0.3, 0.5]$. (b) Effect of distance penalty coefficient $\lambda$, showing that prioritizing left-positioned tokens improves accuracy with minimal speed cost. (c) Comparison of Streaming Parallel Decoding versus block-wise decoding across entropy thresholds; streaming achieves up to $1.9\times$ speedup by enabling immediate prefix commitment.
• Figure 5: Ablation studies. (a) Pareto frontier on GSM8K showing quality-speed trade-offs across hyperparameter configurations; conservative settings achieve 92.3% accuracy at $1.97\times$ speedup while aggressive settings reach $3.2\times$ acceleration. (b) Block size effect during continued pretraining shows stable performance across $B \in \{4, 8, 32\}$. (c) Attention design and model scale: we compare bidirectional attention within blocks (Bi-Attn Block) against our causal design (Our Method) across model sizes; larger models benefit more from causal adaptation, while bidirectional intra-block attention consistently underperforms.