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

TL;DR

WeDLM is a masked diffusion language model that runs entirely under standard causal attention, eliminating the prefix-KV-cache incompatibility that has prevented prior DLLMs (LLaDA, Dream, SDAR) from beating optimized AR engines like vLLM in wall-clock time. The trick is Topological Reordering: observed tokens are physically moved to the front of the sequence while RoPE position ids preserve their logical positions, so masked positions can attend to the full observed context under an unmodified causal mask. Combined with a streaming decoder that biases resolution left-to-right and a dynamic sliding window that refills masks on-the-fly, WeDLM-8B (continued-pretrained from Qwen3-8B on 100B tokens) achieves ~3× speedup over vLLM-served Qwen3-8B on reasoning benchmarks and up to ~10× on low-entropy generation, while matching or exceeding the base model on quality.

Key claims

In KV-cached deployment, the right efficiency metric is prefix cacheability P_c = N_new / N_total — the fraction of processed token instances that end up as cache-reusable committed prefix — not tokens-per-forward [§3.1, Eq. 4].
Bidirectional attention is not required for masked-token recovery; mask recovery only needs each masked position to access all observed tokens, which can be achieved under a strict causal mask [§3.2].
Topological Reordering permutes the physical sequence so observed tokens precede masked tokens, while RoPE position ids preserve logical order — under causal masking each masked token then attends to all observed tokens [§4.1, Eq. 5-6].
Dual-Stream Masking trains a clean memory stream alongside a masked prediction stream that share positional encodings, simulating prefix-conditioned inference and closing a train–inference gap that uniform-masking objectives introduce [§4.2, Eq. 8-10].
Streaming Parallel Decoding combines a distance-penalized confidence rule (Eq. 11) with a fixed-size sliding window that refills committed slots, eliminating the stop-and-wait pipeline bubble of block-wise DLLMs [§5.2].
WeDLM-8B averages 74.72 across 10 reasoning/math/code benchmarks, surpassing Qwen3-8B base (72.61), LLaDA-8B (55.44), and Dream-7B (56.91) [Table 1].
HumanEval climbs from 68.90 (Qwen3-8B) to 75.00 (WeDLM-8B); HumanEval-plus from 63.40 → 68.90; GSM8K from 85.97 → 90.20; MATH from 50.80 → 53.60 [Table 1].
WeDLM-8B served via vLLM achieves ~3× end-to-end speedup on complex reasoning and >10× on low-entropy generation versus the same vLLM-served Qwen3-8B baseline [§1, §6.4].
On a 1-to-200 counting task WeDLM resolves on average 8.10 tokens per forward pass, illustrating throughput gains in low-entropy regimes [project page].
Because attention remains a standard causal mask, each decoding step is a small causal prefill over the active window on top of cached KV — directly compatible with FlashAttention, PagedAttention, and CUDA Graphs without kernel changes [§5.2].

Method

WeDLM reframes masked diffusion decoding to be prefix-cache friendly. During training, given a clean sequence with logical positions, a random subset of tokens is masked; the sequence is then topologically reordered so all observed tokens occupy the physical prefix and all masked tokens occupy the physical suffix, but RoPE position ids still encode the original logical positions. Under an unmodified lower-triangular causal mask, each masked position (physically at the end) therefore attends to every observed token. To bridge the gap between uniform random masking at training time and the suffix-heavy mask distribution at inference, Dual-Stream Masking concatenates a clean memory stream with a block-masked prediction stream sharing position ids; prediction-stream tokens in block i attend to memory-stream tokens whose logical positions precede i, simulating the inference-time setting where earlier blocks are already finalized. An auxiliary next-token AR loss is kept during training to preserve causal LM capability. Initialization is from Qwen2.5-7B / Qwen3-8B with 100B tokens of continued pretraining and 10B tokens of SFT; Magi Attention handles the irregular reordered attention patterns without custom kernels.

At inference, the decoder maintains a fixed-size window of slots mixing already-predicted tokens and [M] placeholders. Each step reorders the window so filled tokens come first, runs a causal forward on top of the persistent KV cache, commits the leftmost contiguous filled prefix (which is immediately cache-valid because of strict causality), picks additional masks to fill using distance-penalized entropy H̃_j = H_j + α·d_j (where d_j is distance from slot j to the leftmost remaining mask), and refills the freed slots with new [M] tokens. The distance penalty biases resolution toward the left, lengthening the contiguous committed prefix and raising P_c; the dynamic refill eliminates the block-boundary synchronization that bottlenecks methods like SDAR and NBDiff.

Results

WeDLM-8B (Qwen3-8B base) averages 74.72 across ARC-C/E, HellaSwag, MMLU, GSM8K, MATH, GPQA-Diamond, MBPP, HumanEval, HumanEval-plus, vs. 72.61 for Qwen3-8B-Instruct, 55.44 for LLaDA-8B, and 56.91 for Dream-7B [Table 1]. Largest wins are on code (HumanEval 68.90 → 75.00; HumanEval-plus 63.40 → 68.90) and math (GSM8K 85.97 → 90.20; GPQA-Diamond 37.00 → 42.42). WeDLM-7B (Qwen2.5-7B base) similarly lifts its base from 67.21 → 70.84 [Table 1]. Wall-clock comparisons are run against vLLM-served Qwen3-8B and report ~3× end-to-end speedups on hard reasoning and >10× on low-entropy generation; on the 1-to-200 counting task WeDLM resolves 8.10 tokens per forward on average. Inference uses window size N=64 with a unified per-step decoding scheme; the distance-penalty coefficient is set to balance quality and speed.

Why it’s interesting

The framing argument — that prefix cacheability dominates over tokens-per-forward in KV-cached serving — reframes the DLLM-vs-AR debate around something concrete and measurable, and explains why so many parallel-decoding papers haven’t translated their per-step token counts into real wall-clock wins. The Topological Reordering + RoPE-decoupling trick is a clean, kernel-free way to graft diffusion-style mask recovery onto an existing AR checkpoint while keeping it deployable on vLLM/FlashAttention/PagedAttention infrastructure. For any team thinking about non-AR or hybrid decoding for generative workloads (including text-conditioned media generators), the “keep the attention mask boring, change the physical permutation” recipe is worth borrowing — it removes the usual “but how do I serve it?” objection that has held back diffusion-LM deployments.