Focus-dLLM: Accelerating Long-Context Diffusion LLM Inference via Confidence-Guided Context Focusing

TL;DR

Diffusion LLMs enable long-context processing but suffer high computational cost from bidirectional full attention. The authors propose Focus-dLLM, a training-free framework that predicts unmasked positions from previous-step confidence (past confidence-guided indicator) and preserves key attention sinks via sink-aware sparse attention with cross-layer reuse, coupled with dynamic KV pruning. Experiments on UltraLLaDA and Dream-7B-Instruct show Focus-dLLM delivers substantial speedups (up to ~29x) while maintaining or improving accuracy, outperforming prior acceleration methods on long-context benchmarks. This work makes long-context diffusion LLMs more practical for real-world, long-input tasks by significantly reducing inference overhead without sacrificing quality.

Abstract

Diffusion Large Language Models (dLLMs) deliver strong long-context processing capability in a non-autoregressive decoding paradigm. However, the considerable computational cost of bidirectional full attention limits the inference efficiency. Although sparse attention is promising, existing methods remain ineffective. This stems from the need to estimate attention importance for tokens yet to be decoded, while the unmasked token positions are unknown during diffusion. In this paper, we present Focus-dLLM, a novel training-free attention sparsification framework tailored for accurate and efficient long-context dLLM inference. Based on the finding that token confidence strongly correlates across adjacent steps, we first design a past confidence-guided indicator to predict unmasked regions. Built upon this, we propose a sink-aware pruning strategy to accurately estimate and remove redundant attention computation, while preserving highly influential attention sinks. To further reduce overhead, this strategy reuses identified sink locations across layers, leveraging the observed cross-layer consistency. Experimental results show that our method offers more than $29\times$ lossless speedup under $32K$ context length. The code is publicly available at: https://github.com/Longxmas/Focus-dLLM

Figures (8)

• Figure 1: Confidence dynamics analysis for LLaDA-8B-Instruct nie2025llada on GSM$8K$cobbe2021gsm8k ($L=76$, $N=128$, and $T=128$). (Left) Confidence score correlation between adjacent steps. (Right) Step-wise recall rates of predicting the unmasked tokens at $t$ using the remasked tokens with top-$4$ highest confidence scores at $t-1$.
• Figure 2: Attention patterns across decoding steps and layers in LLaDA-8B-Instruct nie2025llada ($L=49$, $N=128$, $T=128$). More visual results can be found in the Appendix.
• Figure 3: Overview of Focus-dLLM. We predict unmasked positions at the current step using previous confidence scores. These positions act as queries to retrieve relevant prompt blocks, where attention is computed over the union of these blocks and dynamically identified attention sinks.
• Figure 4: Niah kamradt2023niah results on UltraLLaDA he2025ultralladascalingcontextlength under long-context settings with a maximum context length of $32K$ across different layer depths.
• Figure 5: Efficiency evaluation. Comparison of decoding throughput (tokens/s) on UltraLLaDA he2025ultralladascalingcontextlength (Left) and Dream-7B-Instruct ye2025dream7bdiffusionlarge (Right) across varying context lengths. Red numbers indicate the speedup ratio of Focus-dLLM relative to the Vanilla baseline.