Retrofitting Linear Attention into Diffusion Language Models

2026-08-06

AI Summary

Diffusion language models incur substantial computational overhead during inference due to full attention over historical context, which limits parallel decoding efficiency. This work proposes a block-mixed attention mechanism that retains precise softmax attention within the current denoising block while employing linear attention for historical blocks. Coupled with the LoLCAT post-training strategy, this approach efficiently integrates the hybrid attention into pretrained models. To the best of our knowledge, this is the first effort to enable efficient linear attention in pretrained diffusion language models, achieving significant throughput gains by replacing only a subset of attention layers. Experiments show that LLaDA-Hybrid maintains performance close to the original model on HumanEval, MBPP+, and CMATH benchmarks while achieving up to 1.7× higher inference throughput and supporting greater request concurrency.

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.

Problem

• diffusion language models
• inference acceleration
• prefix attention
• linear attention
• blockwise decoding

Innovation

• block-hybrid attention
• linear attention
• diffusion language models
• parallel decoding
• KV caching