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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Liu, Aiwei, He, Minghua, Zeng, Shaoxun, Zhang, Sijun, Zhang, Linhao, Wu, Chuhan, Jia, Wei, Liu, Yuan, Zhou, Xiao, Zhou, Jie
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866911341343670272
author Liu, Aiwei
He, Minghua
Zeng, Shaoxun
Zhang, Sijun
Zhang, Linhao
Wu, Chuhan
Jia, Wei
Liu, Yuan
Zhou, Xiao
Zhou, Jie
author_facet Liu, Aiwei
He, Minghua
Zeng, Shaoxun
Zhang, Sijun
Zhang, Linhao
Wu, Chuhan
Jia, Wei
Liu, Yuan
Zhou, Xiao
Zhou, Jie
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2512_22737
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference
Liu, Aiwei
He, Minghua
Zeng, Shaoxun
Zhang, Sijun
Zhang, Linhao
Wu, Chuhan
Jia, Wei
Liu, Yuan
Zhou, Xiao
Zhou, Jie
Computation and Language
68T50
I.2.7
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.
title WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference
topic Computation and Language
68T50
I.2.7
url https://arxiv.org/abs/2512.22737