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

Summary

WeDLM argues that the reason diffusion language models have not delivered wall-clock speedups over industrially served autoregressive LLMs is not per-step parallelism but prefix cacheability: how quickly proposed tokens can be committed into a growing left-to-right prefix whose KV states never need recomputation. The paper formalises this intuition with a $p_{cache}$ metric and then makes a clean architectural move — perform masked-diffusion recovery entirely under standard causal attention, with observed tokens physically shifted to the prefix (via Topological Reordering) while RoPE preserves logical positions. This is complemented by Dual-Stream Masking to align training with the suffix-heavy distribution of streaming inference, and by a Streaming Parallel Decoding algorithm that biases resolution toward earlier positions and continuously refills a sliding window. The conceptual contribution is real: reframing DLLM inference around prefix cacheability, and demonstrating that mask-recovery does not require bidirectional attention, is an elegant reconciliation between diffusion training and the mature AR serving stack (vLLM, FlashAttention, PagedAttention). Empirically, WeDLM-7B/8B initialised from Qwen2.5-7B / Qwen3-8B are reported to match or exceed the AR checkpoints on average while delivering ~3× end-to-end speedups over vLLM-served AR baselines and much larger speedups on low-entropy generation. The main conceptual limitation is that the paper's central efficiency claim leans on a batched-throughput comparison whose vLLM baseline runs suspiciously slowly (~110 tok/s at BS=1, ~258 tok/s at BS=16 for Qwen3-8B), the exact deployment settings of which are not documented; the same claim rests on tables (3–5) whose supporting figures and per-baseline instrumentation are not shipped with the released code, and whose numbers include a $p_{cache}$ upper bound and a "tok/fwd > window size" metric that appear internally inconsistent. The main-tables quality claims are also produced under a step-wise one-token-per-position decoder, which measures the model's underlying capability rather than the deployed system's quality at the reported speeds. Overall this is a strong conceptual contribution that would benefit from tighter empirical bookkeeping.

Strengths

• Conceptual contribution. The reframing of DLLM efficiency around prefix cacheability (rather than tokens-per-forward) is a genuine step forward: it identifies why prior DLLMs failed to beat vLLM and pinpoints the fix. The $p_{cache}$ metric, while under-specified, gives a useful vocabulary that the field has been missing.
• Architectural elegance. Topological Reordering is a clean idea. It leaves the causal mask untouched, integrates immediately with FlashAttention/PagedAttention/CUDA Graphs, and requires no custom kernel — the ecosystem-compatibility story is credible and matters practically.
• Initialisation from strong AR checkpoints. Rather than a from-scratch DLLM, WeDLM adapts Qwen2.5-7B / Qwen3-8B with 100B tokens of continued pretraining and gains most of its quality from the AR backbone. This is a pragmatic recipe that others can replicate on top of existing open-weight releases.
• Streaming decoding vs block-wise stop-and-wait. The streaming algorithm with a sliding window and immediate prefix commitment is a coherent operationalisation of the causal-attention design and is empirically shown to beat block-wise decoding on the same model.
• Honest disclosure of variance. The case studies span low-, medium- and high-entropy prompts (8× throughput spread) and the paper acknowledges the high-entropy regime as a genuine limitation, rather than only reporting peak numbers.

Weaknesses

• Vs-AR baseline speed is under-specified and looks under-tuned. In Tables 3–4 the vLLM-served Qwen3-8B baseline reports only ~110 tok/s at BS=1 and ~258 tok/s at BS=16 on math. These are low compared to published vLLM numbers on H20 and scale near-linearly then flatten unusually; the paper does not document the vLLM version, prefill/chunked-prefill settings, input/output length distribution, CUDA graph state, or warmup. Given the paper's headline claim ("first DLLM to outperform industrial AR engines"), the fairness of the comparison rests entirely on this configuration.
• Main-tables quality decoded one-token-per-step, not at the speedup settings. Tables 1 and 2 use "each step generates one token by selecting the lowest-entropy position," which effectively disables the parallelism that motivates the whole paper. So the quality tables measure model capability, not deployed-system quality at the reported speeds. Only Figure 5(a) partially closes this gap, and only on GSM8K.
• Regressions elided from the narrative. Table 1 shows HellaSwag drops by ~3 points for both WeDLM-7B and WeDLM-8B, and MATH drops 4.8 points on WeDLM-8B, but the Section 5.3 summary emphasises only the wins ("MMLU +3.6, GPQA +3.5, HumanEval +8.5"). MBPP is acknowledged as the only consistent regression but never analysed, despite HumanEval — a sibling code benchmark — being one of the strongest gains.
• Statistical rigour. Every reported number is a single run at $T=0.1$ with no seeds documented and no error bars in any ablation. Several headline deltas (MMLU +3.6, GPQA +3.5, block-size 0.8-point spread, $\lambda$ +2.6-point improvement) sit near typical per-seed noise for their benchmarks.
• $p_{cache}$ upper bound is asserted, not proved. The claim that "the theoretical upper bound for $p_{cache}$ in current DLLMs is 0.5" is stated as a hard bound but is really a property of a particular decoding recipe; WeDLM's dual-stream training uses each token twice (memory + prediction), matching the pattern the bound calls unavoidable. The subsequent claim that 0.45–0.48 "approaches the optimal limit" therefore either needs formal argument or should be reframed as empirical.
• DLM tok/fwd exceeds window size. Table 3 reports Math tok/fwd of 6.58 with the paper's $W=6$ window, but "committed per forward" cannot exceed the window. Either the batched throughput runs use a different $W$ (in which case the setting must be stated), or the metric is mis-labelled.
• Position-aware confidence formula elided. $\tilde{H}_i$ is introduced, referenced as "Eq. 9," and never actually shown in the OCR. Since the streaming decoder centrally depends on the sign and functional form (is it $H_i - \lambda d_i$ or $H_i + \lambda d_i$?), this is a substantive gap.
• Scale conclusion overreaches. "7B+ recommended" is drawn from three points (0.6B, 1.5B, 7B) with monotone improvement. There is no evidence about behaviour at 3B–4B or 13B, and no head-to-head with published block-diffusion competitors like Efficient-DLM or SDAR at matched settings.
• Local inconsistencies. WeDLM-8B average 77.36 (intro) vs 77.53 (Table 1 / §5.3); Figure 4(c) caption "11.9×" vs body "1.9×"; Figure 4(a) $\tau$ range [0.3, 0.5] vs body [0.3, 0.6]; Figure 8 caption "Low Entropy" vs body "Medium Entropy"; baseline $p_{cache}$ range quoted as 0.08–0.13 while Table 5 tops out at 0.12; the "~10× reduction in training tokens" phrasing does not match the paper's own derived 25× and 12.5× figures.
• Ablation controls under-specified. The Dual-Stream ablation compares against uniform whole-sequence masking, conflating two design choices (dual-stream construction and block-suffix distribution). The Streaming vs block-wise comparison uses only WeDLM as its own baseline rather than a published block-diffusion system.

Reproducibility & code

• Instruct/base quality tables. The evaluation harness (evaluation/wedlm_eval.py, evaluation_base.sh) and the AR vLLM harness (evaluation/ar_baseline_eval.py) cover the 10-benchmark suite, and WeDLM-7B/8B weights are on HuggingFace. So the WeDLM and Qwen columns of Tables 1 and 2 are code-reachable, modulo the unset seed at $T=0.1$ and small hardware/vLLM drift. The 77.36 vs 77.53 headline mismatch means a reproducer should focus on the sign/magnitude of the ~2-point gap, not the exact value.
• DLLM baseline rows are not reachable from the release. LLaDA-8B, Dream-7B, and SDAR-8B are served under third-party engines (dInfer, JetEngine) with no integration, configuration or documentation shipped. Both the Tables 1/2 DLLM-baseline columns and the Table 5 $p_{cache}$ baseline rows therefore cannot be reproduced from this repository.
• Batched throughput and $p_{cache}$ (Tables 3–5). The engine records decode_tokens and total generation time so raw throughput is measurable, but there is no ready-made script that segments by task (Math/Code/QA), sweeps batch size {1, 8, 16}, or aligns AR (vLLM) and DLM measurements under matched conditions. $p_{cache}$ instrumentation for LLaDA/Dream/SDAR is absent. The reproducer must both build the sweep and pin the vLLM configuration.
• Streaming vs block-wise (Figure 4c). The released engine implements the streaming variant; there is no CLI switch that runs the same model under a block-committed decoder. Setting --wedlm-entropy-threshold None yields one-by-one decoding, which is not equivalent. Reproducing this ablation requires implementing the block-wise baseline.
• Training-time ablations (Figures 5(b), 5(c), 6(a), 6(b)). Block-size sweep, Bi-Attn Block vs AR-continue vs WeDLM at multiple scales, dual-stream vs uniform masking, and aux-AR-loss on/off all require re-running continued pretraining on the 100B-token corpus. That corpus is not provided; the shipped finetune/ code is a "reference implementation," not the internal pipeline; and specifically the Bi-Attn Block, AR-continue and uniform-masking baselines are not exposed as configurable modes. enable_ar_loss is a config flag, but no "no aux-AR" checkpoint is released. In short, essentially every architectural ablation is out of reach for an external reader.
• Case-study tok/s (Figures 7–9). example.py prints tokens, time and tok/s for a single prompt, so the qualitative low/medium/high entropy ordering is reproducible. The absolute numbers (1673.3, 745.2, 197.8) are single-shot, hardware-specific peaks (paper cites H20) and forward-count is not exposed by default.
• Sampling seed. Generative benchmarks run at $T=0.1$ without a documented seed in either paper or code, so exact numbers cannot be recovered even on matched hardware.

Recommended Changes

Essential

• Fully document the vLLM baseline. Given the "first DLLM to outperform industrial AR" headline, add an appendix section stating the vLLM version, all serving flags (chunked prefill, --enforce-eager, prefix caching, KV cache size), input/output length distributions, warmup protocol, and hardware. Explain why Qwen3-8B on H20 delivers only ~258 tok/s at BS=16 in Table 3, or re-run against a properly tuned vLLM. This addresses the "Vs-AR baseline speed is under-specified" weakness.
• Report deployed-quality tables at the reported speed settings. Complement Tables 1 and 2 with a companion table showing quality at $\tau \in \{0.3, 0.5, 0.9\}$ and $\lambda = 0.10$ (the settings that produce the speed numbers), not just the one-token-per-step "unified" scheme. Extend the Figure 5(a) Pareto analysis beyond GSM8K to at least MATH, HumanEval and MMLU. Addresses "Main-tables quality decoded one-token-per-step, not at the speedup settings."
• Add error bars and seeded runs. Set an explicit seed in evaluation/wedlm_eval.py and ar_baseline_eval.py, and report mean ± std over $\geq 3$ seeds for the temperature-> 0 benchmarks (GSM8K, MATH, GPQA, MBPP, HumanEval). Addresses "Statistical rigour" and the "Sampling seed" reproducibility gap.
• Prove or reframe the $p_{cache}$ upper bound. Either supply a formal derivation showing that any DLLM subject to the paper's mask-recovery objective has $p_{cache} \leq 0.5$, or reframe the 0.45–0.48 numbers as strong empirical performance rather than "approaching the optimal limit." Addresses the "$p_{cache}$ upper bound is asserted, not proved" weakness.
• Fix the "tok/fwd exceeds window" contradiction. Either state the actual window size used in Table 3 (if larger than $W=6$) or relabel the metric to "predictions per forward" (or similar), so the number is compatible with the algorithm. Addresses "DLM tok/fwd exceeds window size."
• Release the block-wise decoding baseline and the DLLM-baseline glue. Add --decoder=blockwise to the released engine so Figure 4(c) is reproducible with WeDLM, and add a baselines/ directory with pinned dInfer/JetEngine setup for LLaDA/Dream/SDAR (Tables 1, 2, 5). Addresses the "DLLM baseline rows are not reachable" and "Streaming vs block-wise" reproducibility items.
• Write the missing distance-adjusted entropy equation. Typeset $\tilde{H}_i$ explicitly (functional form and sign) where "We define a distance-adjusted entropy:" currently lands. Addresses "Position-aware confidence formula elided."

Suggested

• Give MATH and HellaSwag regressions equal narrative weight. In §5.3 and the abstract, present a symmetric account of gains and losses (HellaSwag $-3$ pts, MATH $-4.8$ pts on 8B) rather than listing only wins. Analyse why MBPP consistently regresses while HumanEval consistently improves. Addresses "Regressions elided from the narrative."
• Add a mid-scale point (3B–4B) or an upper-scale point (13B/14B). The "7B+ recommended" claim needs more than three monotone points. Addresses "Scale conclusion overreaches."
• Include a head-to-head vs Efficient-DLM. Direct accuracy and tok/s comparison against Efficient-DLM (identified as the closest concurrent competitor) on 2–3 benchmarks under matched hardware would materially strengthen the empirical case. Addresses the "Positioning vs Efficient-DLM" weakness.
• Cleaner Dual-Stream ablation. Separately toggle off the dual-stream construction (single-stream + block-suffix masking) and the block-suffix distribution (dual-stream + uniform masking), rather than the current uniform-masking joint ablation. Addresses "Ablation controls under-specified."
• Ship a throughput sweep script. Provide throughput_sweep.sh producing Tables 3 and 4, and a $p_{cache}$-computing script that at minimum covers the WeDLM row of Table 5. Addresses "Batched throughput and $p_{cache}$ (Tables 3–5)" reproducibility item.
• Justify or re-tune the default $\lambda$. $\lambda = 0.10$ is used for main tables but the sweep only covers 0.02–0.08; either extend the sweep to include 0.10, or justify the choice. Addresses "Position-aware confidence uses $\lambda = 0.10$ for main tables but ablated over 0.01–0.05."
• Reconcile local numeric/label inconsistencies. In one editing pass, align: intro average 77.36 with Table 1 (77.53); Figure 4(c) caption 11.9× with body 1.9×; Figure 4(a) $\tau$ range with body; Figure 8 "Low"/"Medium" entropy; baseline $p_{cache}$ range 0.08–0.13 with the Table 5 max of 0.12; and either replace the "~10× reduction in training tokens" phrasing with a figure matching the 4%/8% derivation or drop the specific multiplier. Addresses the "Local inconsistencies" weakness.
• Add hardware and forward-count to case-study captions. State the GPU (H20) and, ideally, mean of 3 runs for Figures 7–9. Expose forward count in example.py. Addresses the "Case-study tok/s" reproducibility item.