Source-linked AI summary
Focus-dLLM: Accelerating Long-Context Diffusion LLM Inference via Confidence-Guided Context Focusing
Lingkun Long, Yushi Huang, Shihao Bai, Ruihao Gong, Jun Zhang, Ao Zhou, Jianlei Yang
TL;DR
Long-context dLLM inference is costly because bidirectional full attention must operate while future unmasked positions remain unknown, limiting existing sparse-attention methods. Focus-dLLM predicts those positions from past confidence and prunes attention with sink preservation and cross-layer reuse. It reports over 29× speedup at 32K context length while maintaining performance against state-of-the-art baselines.
Problem
Sparse attention for long-context dLLMs must estimate importance for tokens whose future unmasked positions are unknown, and existing coarse estimates limit efficiency.
Method
Focus-dLLM is a training-free framework that predicts unmasked regions from previous-step confidence and applies sink-aware pruning with cross-layer sink reuse.
Results
Over 29× speedup is achieved at 32K context length while maintaining superior performance compared to state-of-the-art baselines.
Takeaways & Limitations
Focus-dLLM removes redundant long-context attention computation during dynamic dLLM decoding without compromising reported performance.
Takeaways & Limitations
The method is evaluated on text tasks, and manually configured hyperparameters may not be optimal across specialized domains.
Abstract
from arXiv · showhide
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
1 Introduction
dLLMs enable parallel long-context generation but remain expensive because diffusion decoding uses bidirectional full attention. Focus-dLLM addresses unknown future unmasked positions with confidence-guided prediction and sink-aware sparse attention.
- Motivation: dLLMs generate tokens through parallel updates with bidirectional attention, supporting throughput and long-context scaling.They replace left-to-right emission with iterative denoising over fixed-length sequences.
- Motivation: Existing efficient-inference methods remain limited because approximate KV caches still attend over full contexts, while sparse methods struggle to estimate future token importance.Unmasked positions are unknown in advance during diffusion, making coarse estimates inaccurate and efficiency suboptimal.
- Motivation: Adjacent-step confidence correlation enables prediction of current unmasked positions from previous-step confidence scores.This finding motivates using past confidence to locate tokens likely to be decoded next.
- Motivation: Cross-layer consistency of attention sinks enables identifying sink locations at an intermediate depth and reusing them in deeper layers.This avoids repeatedly re-identifying influential sink tokens.
- Approach: Focus-dLLM is a training-free framework combining a past confidence-guided indicator with sink-aware dynamic token pruning.It predicts unmasked positions, preserves important sinks, and computes attention only over predicted queries and selected key-value pairs.
- Results: Focus-dLLM preserves accuracy while providing substantial speedups, including 2.05× over Fast-dLLM for UltraLLaDA at 32K context length.The reported comparison also describes better-than-vanilla performance.
2 Related Work
dLLMs are presented as a non-autoregressive text-generation paradigm based on iterative denoising, with recent masked diffusion models scaling to billion-parameter sizes and competitive performance.
- Diffusion large language models: dLLMs generate text through parallel token generation via iterative denoising rather than sequential autoregressive emission.Prior work includes both continuous-space and discrete-token diffusion formulations.
- Diffusion large language models: Recent masked diffusion LMs have scaled to billion-parameter models and demonstrated competitive performance against autoregressive counterparts.The supplied passage frames long-context capability as an active direction within this literature.
3 Preliminaries
dLLM inference iteratively denoises a fixed-length sequence initialized with a prompt followed by masks. At each step, confidence determines which masked positions are unmasked while others are remasked.
- Diffusion LLM inference: Unlike autoregressive models, dLLMs generate text by iteratively denoising a fixed-length sequence.The sequence is processed through repeated updates rather than sequential token emission.
- Initialization: Given a prompt of length M, inference initializes a length-L sequence by appending N = L − M mask tokens.The initial masked positions are M(0) = {M + 1, . . . , L}.
- Denoising updates: At each denoising step, the model predicts a conditional distribution and confidence score for every masked position.The confidence-driven strategy uses these predictions to determine subsequent sequence updates.
- Denoising updates: The decoding strategy unmasks the highest-confidence positions while remasking the remaining positions.This produces the dynamic masked-position process used during diffusion inference.
- Approximate KV cache: Approximate KV caching refreshes selected token states while reusing cached states for the remaining tokens.This reduces computation costs, although bidirectional attention makes conventional KV caching inapplicable.
4 Motivation
Focus-dLLM is motivated by two regularities in dLLM inference: confidence predicts which tokens will be unmasked next, and attention sinks show consistency across layers.
- Motivation: The motivation study examines token-confidence consistency and attention patterns as foundations for Focus-dLLM.These analyses target redundancy in attention computation for tokens likely to be unmasked first.
- Temporal Consistency of Confidence: 96.1% average recall is obtained when top-4 remasked tokens from the previous step predict tokens decoded at the next step.The result is measured across decoding steps in the confidence dynamics analysis.
- Temporal Consistency of Confidence: Tokens unmasked at step t can be reliably located using the confidence of tokens at the preceding step t − 1.This supports predicting active query positions before the current decoding step.
- Spatial Consistency of Attention Sinks: Attention maps exhibit locality, while attention sinks emerge and evolve across denoising steps.The observed sink dynamics motivate preserving these influential tokens during sparse attention.
- Spatial Consistency of Attention Sinks: Attention-sink indices typically match across layers, enabling identification at an intermediate depth and reuse in deeper layers.This cross-layer consistency is presented as a way to eliminate redundant computation.
5 Focus-dLLM
Focus-dLLM combines previous-step confidence prediction, local query expansion, sink-aware token pruning, and KV caching to reduce redundant long-context dLLM attention. It preserves dense attention in initial layers, reuses identified sink locations in later layers, and selects relevant prompt blocks for sparse attention.
- 5 Focus-dLLM: The inference workflow combines KV caching with sparse attention during denoising steps, while performing a full cache refresh at each block entry.Within each block, unmasking is non-autoregressive, while blocks are processed autoregressively from left to right.
- 5.2 Past Confidence-Guided Indicator: Focus-dLLM uses previous-step confidence scores to predict likely future unmasked positions and form focused query tokens.Among positions still masked at step t, the method selects top-k candidates using confidence from step t−1, with k determined by the prediction expansion factor ρ.
- 5.2 Past Confidence-Guided Indicator: The predicted query positions are expanded into local windows so attention retains nearby semantic context while excluding distant tokens.The window size is w, and the active position set is the union of windows around the predicted positions.
- 5.3 Sink-Aware Sparse Attention: Focus-dLLM retains attention sinks because existing sparse dLLM methods may discard tokens pivotal for generation quality.The method identifies sinks dynamically and shares their locations across layers using cross-layer consistency.
- 5.3 Sink-Aware Sparse Attention: The first l_dense layers use full attention, after which sink locations identified at the cut-off layer are reused in sparse layers.The sink set is formed by selecting the top-N_sink tokens according to aggregated attention-importance scores across heads.
- 5.3 Sink-Aware Sparse Attention: Prompt tokens are partitioned into blocks, ranked by relevance to predicted queries, and pruned by retaining the top C blocks.Each block uses a mean key representative, and the selected prompt tokens are combined with response tokens for sparse attention.
6 Experiments
Experiments on UltraLLaDA and Dream-7B-Instruct evaluate Focus-dLLM for accuracy, retrieval, throughput, and accuracy–efficiency trade-offs across long context lengths. Focus-dLLM preserves strong performance while delivering increasing speedups as context length grows.
- Focus-dLLM achieves the highest average score on UltraLLaDA, outperforming Vanilla and all competing acceleration frameworks.
- Focus-dLLM performs on par with Vanilla on Dream-7B-Instruct, surpasses Sparse-dLLM and Fast-dLLM, and reaches up to a 19.95× speedup at 32K context length.
- Focus-dLLM achieves higher Niah scores than Fast-dLLM and Sparse-dLLM across layers and exceeds Vanilla accuracy at the deepest layer.
- 29.6× speedup over Vanilla at 32K context length, increasing from 9.4× at 8K as redundant attention becomes more costly for longer sequences.
- At 16K context length, Focus-dLLM forms a stronger accuracy–throughput Pareto frontier than prior approaches.
7 Ablation Study
Ablations show that both confidence-guided query selection and sink-aware sparse attention contribute to Focus-dLLM’s accuracy and throughput. Hyperparameter results further indicate that retention, prediction coverage, depth, and window size require balanced settings.
- Combining PCGI and SA Sparse Attn achieves further accuracy gains and the highest throughput over the Fast-dLLM baseline.
- SA Sparse Attn improves accuracy and significantly increases throughput, whereas PCGI alone slightly degrades accuracy.
- Adding attention sinks consistently improves LongBench accuracy across Dream-7B-Instruct tasks.
- Increasing sparsity ratio α generally improves accuracy, but Dream performance drops at α=0.7 when excessive retention introduces irrelevant context.
- Larger prediction expansion factors improve performance after small values cause poor accuracy from insufficient recall of future decoded positions.
- Moderate window sizes improve accuracy, while overly small windows miss local context and excessively large windows introduce unrelated tokens.
8 Conclusion
Focus-dLLM is a training-free framework that accelerates long-context dLLM inference by predicting queries with past confidence and pruning attention while retaining critical history. Experiments report over 29× speedup at 32K context length with superior performance compared to state-of-the-art baselines.
- Focus-dLLM combines a past confidence-guided indicator for query prediction with sink-aware pruning to eliminate redundant computation.
- Over 29× speedup at 32K context length is achieved while maintaining superior performance compared to state-of-the-art baselines.
Limitations
The framework is evaluated primarily on text tasks, while multimodal reasoning remains unexplored and manually configured hyperparameters may not transfer optimally across specialized domains.
- Focus-dLLM’s extension to multimodal reasoning remains a direction for future exploration.
- Manually configured hyperparameters may not achieve optimal performance across all specialized domains.
A Implementation Details
Focus-dLLM combines semi-autoregressive remasking, approximate KV caching, and sink-aware sparse attention with dynamic token selection and cache updates. The implementation uses specialized GPU kernels and preserves dense computation in selected stages or layers.
- Implementation: The method uses specialized Triton GPU kernels for dynamic sparse attention while retaining dense attention in initial layers and during full-cache refreshes.
- Inference procedure: Focus-dLLM follows semi-autoregressive remasking, where a transfer scheduler selects tokens to unmask at each step using confidence scores.
- Evaluation setup: Experiments compare Focus-dLLM with vanilla diffusion inference and acceleration frameworks under a common semi-autoregressive remasking strategy with block length 32.
- Inference procedure: The inference algorithm initializes empty KV caches, predicts dynamic unmasking counts, selects top-k focus indices, and performs full cache refreshes at block entries.
- Sparse attention: Sparse attention computes block relevance from focus queries and cached keys, selects a fraction of blocks, and updates active KV-cache entries.
- Baselines: The framework reuses cached activations or critical KV entries in comparison methods, including Fast-dLLM, Sparse-dLLM, and SparseD’s precomputed sparse patterns.
B Details of Accuracy vs. Efficiency Experiments
This section details the accuracy–efficiency evaluation on LongBench and supplements it with visual analysis of attention locality and sink consistency across layers and denoising steps.
- Experimental settings: Sparse-dLLM configurations vary the retention ratio r, determining the percentage of preserved KV-cache entries.This provides a comparison across different retention settings.
- Experimental settings: SparseD configurations differ in the skip ratio for initial full-attention steps and the selection ratio r.The skip ratio specifies the initial portion of steps using full attention.
- Experimental settings: Focus-dLLM configurations vary the sparsity ratio α, which controls the amount of prompt context retained for attention computation.Other hyperparameters remain consistent with the main-text setup described in section 6.1.
- Attention-pattern analysis: Figure 8 visualizes strong diagonal locality and bright vertical attention sinks across layers and denoising steps.The prominent sink locations show strong positional consistency across layers within the same step.
- Performance comparison: Table 5 compares LongBench performance and throughput for UltraLLaDA baselines and multiple Focus-dLLM configurations.The table also includes Sparse-dLLM and SparseD configurations.