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

1. Introduction

• Prior diffusion models’ bidirectional attention prevents standard KV caching, while out-of-order resolution and block finalization limit reusable left-to-right prefixes.This motivates restoring strict causal structure for cache-friendly diffusion decoding.Source: §1 · p. 2
• WeDLM performs diffusion-style mask recovery under standard causal attention, using Topological Reordering to expose all observed tokens to masked positions.Observed tokens move to the physical prefix while retaining their logical positions, allowing masked positions to access the full observed context under a causal mask.Source: §1 · pp. 2–3
• Streaming Parallel Decoding prioritizes earlier unresolved positions and refills a sliding window as tokens are committed, avoiding block-wise stop-and-wait decoding.Strict causal attention makes newly committed prefix tokens immediately cache-valid.Source: §1 · p. 3
• The framework is instantiated on Qwen2.5-7B and Qwen3-8B, with 100B tokens of continued training and 10B tokens of supervised fine-tuning.Across code, math, and general-knowledge benchmarks, WeDLM preserves or often improves its base models’ capabilities.Source: §1 · p. 3
• WeDLM-8B averages 77.36 versus Qwen3-8B-Instruct’s 75.12, with up to 3× end-to-end acceleration on complex reasoning and over 10× on some low-entropy tasks.The comparisons use the vLLM engine under matched, practical inference conditions.Source: §1 · p. 3
• WeDLM surpasses optimized vLLM baselines by more than 3× on complex reasoning tasks while maintaining generation quality.Figure 1 reports about 3× speedup over vLLM-optimized Qwen3-8B on GSM8K and shows WeDLM-8B-Instruct matching or surpassing Qwen3-8B-Instruct across several benchmark areas.Source: Figure 1 · p. 2, §1 · p. 3

2.1. Autoregressive Language Modeling

• Autoregressive generation is sequential because each token is conditioned on the previously generated prefix.The model factorizes the sequence into conditional probabilities for successive tokens.Source: §2.1 · p. 4

2.2. Decoupled Positional Representation

• RoPE indexes attention scores by supplied logical positions, allowing tokens to be processed in a different physical order.Inputs are represented as token–position pairs, and the attention mask determines allowed information flow.Source: §2.2 · p. 4

2.3. Masked Diffusion Language Models

• Masked diffusion trains models to reconstruct randomly masked tokens from a corrupted sequence using a weighted cross-entropy objective.The loss factor 1/γ compensates for varying numbers of masked tokens across noise levels.Source: §2.3 · p. 4
• Standard masked diffusion typically uses bidirectional attention, which is incompatible with standard KV caching and mismatches causal representations in pretrained autoregressive models.The paper instead optimizes a denoising objective under strictly causal attention using reordering-based context exposure.Source: §2.3 · p. 4

3.1. Prefix Cacheability (𝑝cache) as an Inference Metric

• In KV-cached decoding, prefix cacheability can matter more than per-step parallelism because predictions outside the committed prefix must be recomputed.The average recomputation factor is 1/pcache; improving cacheability can match or exceed speedups from increasing parallelism.Source: §3.1 · p. 5, §3 · p. 5
• Bidirectional KV coupling and out-of-order resolution reduce cacheability; WeDLM addresses them with causal attention and left-to-right commitment mechanisms.Only tokens whose states depend on earlier context can be reused as a growing prefix.Source: §3.1 · p. 5

3.2. Rethinking the Necessity of Bidirectional Attention

• Causal mask recovery can expose every masked position to all observed tokens while ordering masked-token dependencies causally.This directed factorization lets earlier-resolved masks influence later ones without bidirectional attention.Source: §3.2 · pp. 5–6
• The paper does not claim equivalence to bidirectional-attention MDLMs and treats the sufficiency of directed masked-token dependence as an empirical question.Its design requirement is standard causal attention with full observed-context access for each masked position.Source: §3.2 · p. 6

4. WeDLM Training: Causal Mask Recovery

• The framework targets prefix-cache-compatible computation under a strict causal mask while exposing the full observed set to masked positions.At inference, the memory stream is discarded and decoding uses a standard causal attention mask over a single sequence, without architecture changes.Source: §4 · p. 6, §4.1 · p. 6, §4.2 · p. 9

4.1. Causal Mask Recovery via Topological Reordering

• Topological Reordering places observed tokens before masked tokens physically while preserving their logical positions, letting masked tokens access all observed context under causal attention.The model trains to recover ground-truth tokens at masked positions.Source: §4.1 · pp. 6–8

4.2. Dual-Stream Masking for Training

• Dual-Stream Masking addresses the gap between uniform training masks and suffix-like unresolved positions during prefix-conditioned inference.A clean memory stream supplies preceding-block context, while the prediction stream uses causal attention within its current reordered block rather than noisy predictions from earlier prediction-stream blocks.Source: §4.2 · pp. 8–9

5. WeDLM Inference: Streaming Parallel Decoding

• Streaming Parallel Decoding commits sufficiently confident tokens into a growing prefix while refilling masked positions to maintain a steady workload.The dynamic sliding window is designed to avoid stop-and-wait behavior at block boundaries.Source: §5 · p. 9, §5.2 · p. 11

5.1. Inference Requirements for Streaming Decoding

• The decoder must make predicted tokens immediately cache-valid, prioritize earlier unresolved positions, and avoid stalling at block boundaries.WeDLM uses standard causal attention for immediate KV validity and combines position-aware confidence selection with on-the-fly window refilling.Source: §5.1 · pp. 9–10

5.2. Streaming Parallel Decoding

• At each step, WeDLM reorders filled tokens before masks, runs a causal forward over the window, commits the leftmost contiguous filled prefix, then predicts and refills.The distance penalty biases selection toward earlier mask positions, supporting contiguous prefix growth.Source: §5.2 · pp. 10–11
• Each decoding step uses standard causal attention over the current window conditioned on the cached prefix, enabling deployment on existing inference systems without kernel modification.The passage names FlashAttention, PagedAttention, and CUDA Graphs as supported systems.Source: §5.2 · p. 11

6.1. Training Details

• Training starts from Qwen2.5-7B and Qwen3-8B, uses 100B continued-pretraining tokens and 10K internal instruction-response pairs for SFT.The setup also uses block size B=32 and an auxiliary autoregressive loss during training.Source: §6.1 · p. 12

6.3. Performance Evaluation

• WeDLM-8B averages 77.53 versus Qwen3-8B’s 75.12 on instruct models, while base-model averages also improve and diffusion baselines trail.The reported comparison covers diverse benchmarks; the best diffusion instruct baseline averages 74.22.Source: Table 1 · p. 12, Table 2 · p. 13, §6.3 · pp. 13–14

6.2. Evaluation Setup

• Evaluation spans reasoning, knowledge, commonsense, and code benchmarks, with WeDLM served via vLLM alongside AR baselines and diffusion models using their recommended frameworks.For a unified comparison, Tables 1 and 2 use step-wise decoding; the default window size is W=6 and distance penalty is λ=0.10.Source: §6.2 · pp. 12–13

6.4. Speed Evaluation

• On MATH, accuracy stays around 53–54% for τ≤0.5, and increasing λ from 0.01 to 0.05 improves accuracy by 2.6 points with only 3% speed reduction.The authors recommend τ∈[0.3, 0.6] for balanced operation.Source: §6.4 · pp. 14–15
• At τ=0.9, streaming reaches 1.9× speedup over block-wise decoding (423 vs. 221 tokens/s); GSM8K settings range from 1.97× speedup at 92.3% accuracy to 3.2× acceleration above 79%.Streaming commits contiguous prefix tokens without waiting for a full block; the GSM8K frontier exposes a quality-speed trade-off.Source: Figure 4 · p. 14, §6.4 · p. 15

6.5. Ablation Studies

• Performance remains within a 0.8-point range across B∈{4, 8, 32}; causal intra-block attention scores higher overall and permits per-token caching unlike bidirectional blocks.Smaller block sizes incur higher Magi Attention overhead, while bidirectional blocks delay cache-valid commitments until the whole block resolves.Source: §6.5 · pp. 15–16, Figure 5 · p. 15
• WeDLM’s performance change varies with base-model scale: the 0.6B and 1.5B models decline by 3.9 and 0.6 points, while the 7B model improves by 3.6 points.The authors recommend 7B+ starting models and leave systematic verification of a possible scaling law for future work.Source: §6.5 · p. 16

6.6. Case Study

• Throughput falls as output entropy rises: a counting task reaches 1673.3 tokens/s, structured mathematical reasoning reaches 745.2, and open-ended explanation reaches 197.8.The authors note that high-entropy generation sees diminishing speedup and remains a framework limitation.Source: §6.6 · p. 16

Discrete Diffusion Language Models

• LLaDA was the first to scale masked diffusion to 8B parameters, demonstrating competitive performance with autoregressive models; LLaDA-MoE matched dense-model performance using 1/6 active parameters.Masked diffusion models require approximately 16× more compute to match autoregressive perplexity, while exhibiting similar scaling trends [18].Source: §7 · p. 17

Adapting Autoregressive Models to Diffusion

• AR-to-diffusion adaptation methods include DiffuLLaMA’s shift operation and attention-mask annealing [11], and Dream’s context-adaptive noise rescheduling [30].Dream-Coder extended the approach to code generation and reported sketch-first reasoning patterns [27].Source: §7 · p. 17

Block Diffusion and Inference Acceleration

• Block-diffusion methods vary in training, attention, and decoding, including context-causal gradual block growth in NBDiff [25] and confidence-based truncation in SDAR [5].Other examples include vectorized training in BD3-LM [1], longest-prefix decoding in SDLM [16], and scaling block diffusion to 100B parameters in LLaDA2.0 [3].Source: §7 · p. 17

Permutation and Reordering

• Unlike XLNet’s permutation language modeling [29], WeDLM reorders observed tokens to the physical prefix for inference-time parallel decoding while preserving logical positions and the standard causal mask.The passage identifies RoPE position ids as one way to preserve logical positions.Source: §7 · p. 17

8. Conclusion

• Prefix cacheability, formalized as p_cache, measures the rate at which predictions become a contiguous left-to-right prefix and can be reused.The conclusion argues that out-of-order resolution and bidirectional KV coupling reduce the computation amortized by caching.Source: §8 · p. 17
• Topological Reordering exposes all observed context to masked positions while preserving the strict causal mask, making newly committed tokens immediately cache-valid.Streaming Parallel Decoding favors earlier positions and refills a fixed window to grow the prefix with minimal recomputation.Source: §8 · pp. 17–18
• Under matched, cache-enabled decoding, WeDLM retains and often improves strong AR-backbone capabilities while accelerating inference.The conclusion presents causal diffusion as a route to reconcile parallel generation with efficient cached decoding.Source: §8 · p. 18
• The authors recommend treating prefix cacheability as a first-class design objective: predicting many tokens per iteration helps only insofar as they can become a cache-valid prefix quickly.They frame causal diffusion as a principled route to align multi-token prediction with cached decoding.Source: §8 · p. 18

• 1673.3 tokens/s on deterministic counting from 1 to 200 is the low-entropy example’s decoding speed.Source: Figure 6 · p. 21
• 745.2 tokens/s on a linear-equation derivation illustrates the medium-entropy structured-reasoning case.Source: Figure 7 · p. 22
• 197.8 tokens/s on an open-ended quantum-physics explanation reflects reduced effective parallel block size with high semantic diversity and precise lexical selection.Source: Figure 8 · p. 23