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

TL;DR

Diffusion language models incur high inference cost from full-sequence attention at every denoising step. By analyzing token-level locality, the authors propose Window-Diffusion, a training-free framework with a dual-window inference mechanism and phase-level KV caching that selectively updates active tokens while reusing context representations. This approach achieves substantial speedups (up to 99x with adaptive-length inference) with minimal generation-quality loss and is complementary to other acceleration techniques. The work offers a practical, architecture-agnostic path to faster diffusion-based text generation in pretrained models.

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.

Figures (6)

• Figure 1: Illustration of mask-based DLM inference.
• Figure 2: Token-wise prediction confidence during diffusion inference. The heatmaps visualize prediction confidence over undecoded positions at diffusion steps $t=64, 128,$ and $192$ on LLaDA. The total sequence length is 256. Decoded tokens are masked out, and color indicates confidence.
• Figure 3: KL divergence between active-token predictions under truncated undecoded context and the full-sequence, no-cache reference on LLaDA. Results are shown for standard recomputation (No-cache) and for reusing KV representations of non-active undecoded tokens from the previous step (Cache). KL is averaged over active tokens, with shaded bands indicating the 25th--75th percentiles across observation steps (30--60).
• Figure 4: Temporal stability of decoded-token Value representations on LLaDA. (a) Recently decoded tokens: we track tokens at positions 100--120 and compute the average cosine similarity of their Value representations over the first 64 steps after each token is decoded. (b) Earlier-decoded tokens: at an observation step $t_0=64$, we select the first 16 already decoded tokens (excluding the prompt) and measure the average Value similarity across the subsequent 64 diffusion steps.
• Figure 5: Window-Diffusion inference framework.Top: Dual-window token organization within a phase, where only active tokens are fully computed, buffer tokens reuse KV states, and far-field tokens are pruned; the window shifts only at phase boundaries. Bottom: Phase-level KV caching, with a refresh step followed by multiple normal steps that reuse cached KV states until the next phase.