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

LLM PTSAD MQIA
2025年12月28日
自回归(AR)生成是大语言模型(LLM)的标准解码范式,但其逐个生成token的特性限制了推理阶段的并行性。扩散语言模型(DLLM)通过每步恢复多个被掩码的token实现并行解码;然而在实践中,它们往往难以将这种并行性转化为相比优化后的自回归引擎(例如vLLM)的实际部署速度优势。一个关键原因在于,许多DLLM依赖双向注意力机制,这破坏了标准的前缀KV缓存机制,导致必须反复进行上下文重构,从而损害了效率。为此,我们提出了WeDLM——一种完全基于标准因果注意力的扩散解码框架,旨在使并行生成过程对前缀缓存友好。其核心思想是:在保持严格因果掩码的前提下,让每个被掩码的位置都能基于所有当前已观测到的token进行条件建模,这一目标通过“拓扑重排序”(Topological Reordering)实现,即将已观测到的token移至物理上的前缀位置,同时保留其逻辑上的原始位置。基于这一特性,我们进一步提出了一种流式解码流程,持续将高置信度的token加入不断增长的从左到右的前缀序列中,并维持固定的并行计算负载,从而避免了块状扩散方法中常见的“停等”行为。实验表明,WeDLM在保持强自回归主干模型生成质量的同时,实现了显著的速度提升:在具有挑战性的推理任务上接近3倍加速,在低熵生成场景下甚至可达10倍;尤为关键的是,我们的对比是在与vLLM驱动的自回归基线相同部署条件下进行的,证明了扩散式解码在实际应用中确实可以超越优化后的自回归引擎。
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.