Skip to content

Window-Diffusion: Accelerating Diffusion Language Model Inference with Windowed Token Pruning and Caching

Jan 2026 · Advanced Parallel Programming Technologies · pp. 112-125 · 3 citations · 24 references
Computer Science

TL;DR

The token-level analysis reveals pronounced structural locality in DLM inference, and proposes a window-based token pruning and caching method for inference that achieves up to $99\times inference speedup while largely preserving generation performance.

Abstract

Diffusion language models (DLMs) generate text through iterative denoising, but inference requires full-sequence attention at every iteration, resulting in substantial redundant computation on masked tokens. Block-wise diffusion can reduce this cost, yet it typically relies on retraining and constrained update orders, limiting its direct applicability to pretrained DLMs. Our token-level analysis reveals pronounced structural locality in DLM inference. Decoding is driven by a small set of prefix-localized active tokens; the influence of distant undecoded context diminishes rapidly, and decoded tokens exhibit stage-wise temporal stability, enabling reuse of intermediate representations except for a brief post-decode transient. Motivated by these observations, we propose \textbf{\placeholder}\footnote{The source code is available at https://github.com/vhicrgit/Window-Diffusion.}, a window-based token pruning and caching method for inference. We maintain a local computation window that slides rightward as denoising progresses, and partition undecoded tokens into: (i) \textit{active tokens} that are computed online, (ii) \textit{buffer tokens} whose KV states are cached and periodically refreshed, and (iii) \textit{far-field tokens} that are pruned outside the window. Computation is restricted to active and buffer tokens within the window, while far-field tokens are omitted at each stage. Experiments on LLaDA and Dream show that, under matched compute budgets, our method achieves up to $99\times$ inference speedup while largely preserving generation performance.

View source

Similar papers

#artificial intelligence Preprint Oct 2026

ALoDLM: Adaptively Looped Diffusion Language Models

ALoDLM is introduced, which replaces uniform computation with token-adaptive latent recurrence, and retains fast parallel decoding, yielding a strong quality-efficiency trade-off among evaluated autoregressive and diffusion models under optimized inference engines.

Lian-Cheng Fang, Zhuo-Wei Li, Youngeun Kim et al. · 1 citation
#machine learning Preprint Oct 2026

LED: Improving Masked Diffusion Language Models by Lookahead Editing

This design introduces a rewritable proposal that uses additional compute to improve expressivity while keeping the simple mask-based transition schedule of MDMs, and fine-tune Qwen2.5-7B and Qwen3-8B to obtain LED-\texttt{7B} and LED-\texttt{8B}, respectively.

Yu-Yan Zhou, Ze-Yuan Yang, Yuxuan Song et al. · 0 citations
#natural language process... Preprint Oct 2026

Prefilling-dLLM v2: Learned Chunk Routing for Long-Context Diffusion Language Models

Diffusion language models apply bidirectional attention over the whole context at every denoising step, so a longer context makes each step substantially more expensive. Sparsifying only at inference leaves the pretrained weights unchanged, so the model is never trained under the sparse pattern it runs with. We propose...

Qi Han, Jing Xiong, Zi-Yuan Zhang et al. · 0 citations
Preprint Oct 2026

Step-dLLM: Adaptive Step-aware Sparse Attention for Efficient Diffusion LLM Inference

Diffusion large language models (dLLMs) generate tokens in parallel via iterative denoising, but their bidirectional attention prevents the KV caching used by autoregressive models and recomputes the full attention map at every step, incurring quadratic per-step cost. Existing sparse attention methods for dLLMs constru...

Zhi-Chen Zeng, Xi-Chong Zhang, Jun-Pan Wu et al. · 0 citations
#machine learning Preprint Sep 2026

Elastic Threshold Attention: Learned Contextual Sparsity for Long-Context Decoding

This work introduces Elastic Threshold Attention, an end-to-end trainable architecture that rivals dense model quality under hardware-aligned block-sparse decoding, and introduces an offline calibration algorithm for domain-specific deployments that freezes per-head constant thresholds, cutting attention compute by an...

Themistoklis Haris, Henry Li, Maryam Karimzadehgan · 1 citation · ⚡1

Related blog posts

Microsoft Research Blog Nov 3, 2026

Learn the rules: The hacker mindset

Research Manager Mike Walker reflects on a career spanning hacker culture, cybersecurity, and AI, sharing lessons about trust, teamwork, leadership, and why solving difficult problems requires curiosity, judgment, agility—and an understanding of the rules. The post Learn the rules: The hacker mindset appeared first on Microsoft Research.

Google DeepMind Blog Nov 3, 2026

Agentic cheminformatics with Google Antigravity

Learn how Google Antigravity and our Cheminformatics Skills accelerate drug discovery workflows, search compounds, and propose novel rigidified macrocycles.

We use cookies to run the site and, with your consent, for analytics and to show ads. See our Cookie Policy.