Retrofitting Linear Attention into Diffusion Language Models
- URL: http://arxiv.org/abs/2608.06628v1
- Date: Thu, 06 Aug 2026 22:34:20 GMT
- Title: Retrofitting Linear Attention into Diffusion Language Models
- Abstract summary: Diffusion language models (dLLMs) offer a promising alternative to autoregressive models by accelerating inference through parallel decoding.<n>Recent dLLMs commonly use blockwise semi-autoregressive decoding, generating blocks autoregressively while denoising tokens within each active block in parallel.<n>We introduce block-hybrid attention, which retains exact softmax attention within the active denoising block while applying linear attention over previous blocks.
- Score: 8.41282954657158
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: Diffusion language models (dLLMs) offer a promising alternative to autoregressive models by accelerating inference through parallel decoding. Recent dLLMs commonly use blockwise semi-autoregressive decoding, generating blocks autoregressively while denoising tokens within each active block in parallel. However, despite KV caching, each denoising step still attends to all previous blocks, repeatedly incurring prefix-attention cost. Motivated by this bottleneck, we ask whether dLLM inference can be further accelerated by linearizing attention over previous blocks. We introduce block-hybrid attention, which retains exact softmax attention within the active denoising block while applying linear attention over previous blocks. We show that this hybrid attention can be retrofitted into a pretrained dLLM with minimal post-training: LLaDA-Hybrid replaces 6 of the 20 attention layers in LLaDA~2.1, a 16B open-source dLLM, largely following LoLCAT (Zhang et al, 2024). The conversion takes only approximately 60 hours while preserving benchmark performance: 72.0% vs. 75.6% on HumanEval, 63.0% vs. 57.7% on MBPP+, and 86.7% vs. 88.3% on CMATH. With a Triton implementation, LLaDA-Hybrid achieves up to $1.7\times$ higher decoding throughput and supports more concurrent requests before exhausting memory, showing that pretrained dLLMs can be efficiently linearized for faster inference. Our code is available at: https://github.com/Diuven/LLaDA-Hybrid.
Related papers
- Serving Masked Diffusion LLMs: Characterization and Design Principles from Real Hardware [1.7286026533711227]
Masked diffusion language models (dLLMs) can generate text faster than autoregressive (AR) models, since they denoise many tokens at once.<n>Recent systems have begun building serving infrastructure for dLLMs, but none first measure how these models behave under real, concurrent serving.<n>We characterize dLLM serving to close this gap, using LLaDA-8B-Instruct with a D2F (Discrete Diffusion Forcing) LoRA adapter on a single NVIDIA H200 GPU.
arXiv Detail & Related papers (2026-08-24T20:16:48Z) - CAI-DLLM: Convergence Aware Inference for Diffusion Language Models [1.8883742168400344]
Diffusion language models can generate many tokens in parallel, but they still require repeated denoising steps during inference.<n>We propose CAI-DLLM, a training-free inference method that uses first-step confidence to guide denoising and reduce inference time.<n>We evaluate CAI-DLLM on LLaDA-8B-Instruct and Dream-7B-Instruct across math, code, reasoning, commonsense, and long-context tasks.
arXiv Detail & Related papers (2026-08-23T23:03:59Z) - Training Hybrid Block Diffusion Language Models with Partial Bidirectionality [11.618306403714692]
Long-context generation is one of the central challenges for large language models.<n>We propose a BDLM Mamba--attention hybrid that addresses this challenge by restricting the reverse Mamba scan to the active denoising block.<n>For long-context inference, BDLM Mamba-H reaches 19.7x the throughput of full-sequence DiffuMamba-H at 65K tokens and 3.7x the throughput of BDLM attention at 262K.
arXiv Detail & Related papers (2026-07-02T22:37:43Z) - Multi-Block Diffusion Language Models [46.581695254036]
Block Diffusion Language Models (BD-LMs) improve diffusion-based text generation with KV caching and flexible-length generation.<n>We propose Multi-Block Diffusion Language Models (MBD-LMs), obtained by post-training BD-LMs with Multi-block Teacher Forcing (MultiTF)<n>MBD-LLaDA2-Mini increases average Tokens Per Forward pass (TPF) from 3.47 to 6.19 and improves average accuracy from 79.95% to 81.03%.<n>When combined with DMax, MBD-LLaDA2-Mini-DMax reaches an average TPF of 9.34
arXiv Detail & Related papers (2026-06-28T05:53:45Z) - HERALD: High-Throughput Block Diffusion LLM Serving via CPU-GPU Cooperative KV Cache Retrieval [31.594043836825403]
Diffusion LLMs (dLLMs) improve GPU utilization over autoregressive decoding by generating multiple tokens per forward pass.<n>Their KV cache still grows linearly with context, limiting throughput at long contexts.<n>We present HERALD, a KV offloading system for block dLLMs.
arXiv Detail & Related papers (2026-06-19T17:36:51Z) - Roll Out and Roll Back: Diffusion LLMs are Their Own Efficiency Teachers [76.15132587294862]
Wide-In, Narrow-Out (WINO) is a training-free decoding algorithm that enables revokable parallel generation.<n>WINO+ injects the verified denoising trajectories produced by WINO into model parameters, aligning training with efficient inference.<n>Experiments on LLaDA and MMaDA show that WINO improves both quality and efficiency, while WINO+ further strengthens this progression.
arXiv Detail & Related papers (2026-05-16T11:27:40Z) - Prefix-Adaptive Block Diffusion for Efficient Document Recognition [52.15352911463151]
Block Diffusion Models (BDMs) support parallel generation, flexible-length output, and KV caching, making them promising for efficient document parsing.<n>We propose the Prefix- Block Diffusion Model (PA-BDM) which replaces intra-block bidirectional denoising with causal denoising from prefix to suffix.<n> Experiments show that the 3B PA-BDM achieves higher recognition scores on several benchmarks and improves inference throughput by 71.6% over the 2.5B MinerU-Diffusion.
arXiv Detail & Related papers (2026-05-16T07:50:13Z) - DMax: Aggressive Parallel Decoding for dLLMs [77.24184219948337]
We present DMax, a new paradigm for efficient diffusion language models (dLLMs)<n>It mitigates error accumulation in parallel decoding, enabling aggressive decoding while preserving generation quality.<n>At the core of our approach is On-Policy Uniform Training, a novel training strategy that efficiently unifies masked and uniform dLLMs.
arXiv Detail & Related papers (2026-04-09T14:35:42Z) - ES-dLLM: Efficient Inference for Diffusion Large Language Models by Early-Skipping [26.560813832545563]
Diffusion large language models (dLLMs) are emerging as a promising alternative to autoregressive models (ARMs)<n>We analyze the generation dynamics of dLLMs and find that intermediate representations, including key, value, and hidden states, change only subtly across successive iterations.<n>We propose textbfES-dLLM, a training-free inference acceleration framework for dLLM that reduces computation by skipping tokens in early layers based on the estimated importance.
arXiv Detail & Related papers (2026-03-10T14:31:19Z) - dParallel: Learnable Parallel Decoding for dLLMs [77.24184219948337]
Diffusion large language models (dLLMs) offer parallel token prediction and lower inference latency.<n>Existing open-source models still require nearly token-length decoding steps to ensure performance.<n>We introduce dParallel, a simple and effective method that unlocks the inherent parallelism of dLLMs for fast sampling.
arXiv Detail & Related papers (2025-09-30T16:32:52Z) - Fast-dLLM v2: Efficient Block-Diffusion LLM [64.38006546510337]
Fast-dLLM v2 is a block diffusion language model that adapts pretrained AR models into dLLMs for parallel text generation.<n>This represents a 500x reduction in training data compared to full-attention diffusion LLMs such as Dream (580B tokens)
arXiv Detail & Related papers (2025-09-30T14:40:18Z) - Towards Effective and Efficient Non-autoregressive Decoding Using Block-based Attention Mask [74.64216073678617]
AMD performs parallel NAR inference within contiguous blocks of output labels concealed using attention masks.
A beam search algorithm is designed to leverage a dynamic fusion of CTC, AR Decoder, and AMD probabilities.
Experiments on the LibriSpeech-100hr corpus suggest the tripartite Decoder incorporating the AMD module produces a maximum decoding speed-up ratio of 1.73x.
arXiv Detail & Related papers (2024-06-14T13:42:38Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.