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

Lingkun Long^1,3, Yushi Huang^2,3, Shihao Bai³, Ruihao Gong^1,3, Jun Zhang², Ao Zhou¹, Jianlei Yang^1,*

¹ Beihang University, ² Hong Kong University of Science and Technology, ³ SenseTime Research

[📝 Paper] | [🚀 Code]

---

📢 News

• The code of Focus-dLLM is released.
• Focus-dLLM is officially presented.
• Supported Models: LLaDA-8B, UltraLLaDA, Dream-7B.

---

📖 Introduction

Focus-dLLM is a training-free attention sparsification framework tailored for accurate and efficient long-context inference in Diffusion Large Language Models (dLLMs).

While dLLMs introduce a compelling non-autoregressive paradigm via iterative denoising, the considerable computational cost of bidirectional full attention limits their inference efficiency. Existing sparse attention methods remain ineffective for dLLMs because they require estimating attention importance for tokens yet to be decoded, while the unmasked token positions are unknown during the diffusion process.

To address this, Focus-dLLM introduces a novel pipeline that accurately predicts unmasked regions and retains only necessary computation. Our approach delivers more than 29× lossless speedup under 32K context length, offering a superior Pareto frontier between throughput and generation quality.

Figure 1: Overview of Focus-dLLM. We employ a past confidence-guided indicator to predict unmasked positions and leverage a sink-aware pruning strategy to dynamically identify and reuse attention sinks.

---

🛠️ Core Techniques

1. Past Confidence-Guided Indicator

We discover that token confidence exhibits a strong positive correlation across adjacent denoising steps. Based on this, we design an indicator that uses confidence scores from step $t-1$ to accurately predict the unmasked positions at step $t$. We further apply window expansion to these predicted positions to preserve local semantic coherence, forming a focused query set for attention computation.

2. Sink-Aware Sparse Attention

To eliminate redundant computation, we propose a sink-aware pruning strategy. This mechanism:

• Dynamic Attention Sink Identification: Explicitly identifies and retains attention sinks to preserve generation quality.
• Cross-Layer Consistency: Leverages the observation that attention sinks match across layers. We identify sinks at an intermediate layer and reuse their locations for subsequent sparse layers, avoiding repeated re-identification overhead.
• Block-wise Pruning: Applies dynamic pruning to Key/Value states of the prompt tokens to keep the most relevant history while retaining all response tokens.

---

🚀 Key Results

Accuracy

Focus-dLLM achieves robust performance across LongBench. On UltraLLaDA, our method achieves the highest average score, outperforming vanilla baselines and existing acceleration frameworks.

Efficiency

Focus-dLLM demonstrates superior scalability. As context length grows, the speedup ratio expands significantly, reaching 29.6× at 32K context.

---

⚡ Quick Start

git clone https://github.com/Longxmas/Focus-dLLM.git && cd Focus-dLLM
conda create -n focus_dllm python=3.11 -y && conda activate focus_dllm
pip install torch==2.8.0
pip install ./third_party/pyext-0.7 -r focusdllm/requirements/runtime.txt
pip install transformers==4.46.3 "datasets<4" "pyarrow<21" "numpy<2" "setuptools<81" sgl-kernel==0.3.16.post3
export PYTHONPATH=$(pwd)/focusdllm:$PYTHONPATH

import os
import torch
from transformers import AutoTokenizer
from opencompass.models.focus_dllm.llada.modeling_llada_focus import LLaDAModelLM
from opencompass.models.focus_dllm.llada.generate_step_focus import generate_with_focus

os.environ['USE_FUSED_ATTN'] = '0'  # 0: PyTorch SDPA; 1: sgl-kernel FlashAttention (batch size 1)
path = 'relaxe-system-lab/UltraLLaDA'
model = LLaDAModelLM.from_pretrained(path, torch_dtype=torch.bfloat16, trust_remote_code=True).cuda().eval()
tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)

prompt = open('document.txt').read() + '\n\nQuestion: ...\nAnswer:'
input_ids = tokenizer(prompt, return_tensors='pt').input_ids.cuda()
output, nfe = generate_with_focus(model, input_ids, steps=256, gen_length=256, block_length=32)
print(tokenizer.decode(output[0, input_ids.shape[1]:], skip_special_tokens=True))

📊 Evaluation

Set MODEL_PATH and METHODS at the top of each config under focusdllm/myeval/, then run:

cd focusdllm
export TIKTOKEN_CACHE_DIR=$(pwd)/../third_party/tiktoken_cache
export USE_FUSED_ATTN=0

# LongBench
python run.py myeval/eval_long/longbench/eval_llada_long.py --dump-eval-details -r
python run.py myeval/eval_long/longbench/eval_dream_long.py --dump-eval-details -r

# Needle in a haystack
python run.py myeval/eval_long/niah/eval_llada_niah.py --dump-eval-details -r
python run.py myeval/eval_long/niah/eval_dream_niah.py --dump-eval-details -r

# Throughput
# The benchmark selects USE_FUSED_ATTN=1 for Focus-dLLM.
cd myeval/eval_speed
MODEL=llada MODEL_PATH=relaxe-system-lab/UltraLLaDA DATA_PATH=/path/to/LongBench/data/narrativeqa.jsonl bash run_benchmark.sh
MODEL=dream MODEL_PATH=Dream-org/Dream-v0-Instruct-7B DATA_PATH=/path/to/LongBench/data/narrativeqa.jsonl bash run_benchmark.sh

---

🗒️ To-Do List

• Refactor and release the official implementation code.
• Support more dLLM architectures.

---

🤓 Acknowledgement

Our framework is built upon the open-source projects Fast-dLLM and OpenCompass. We also incorporate ideas and implementation details inspired by several pioneering works in the dLLM and sparse attention community:

• Sparse-dLLM
• SparseD

Furthermore, our experiments leverage diffusion language models from:

• UltraLLaDA
• LLaDA
• Dream-7B

We sincerely thank the authors for their open-source contributions, which greatly facilitated the development of Focus-dLLM.

📌 Citation

@article{long2026focusdllm,
  title   = {Focus-dLLM: Accelerating Long-Context Diffusion LLM Inference via Confidence-Guided Context Focusing},
  author  = {Long, Lingkun and Huang, Yushi and Bai, Shihao and Gong, Ruihao and Zhang, Jun and Zhou, Ao and Yang, Jianlei},
  journal = {arXiv preprint arXiv:2602.02159},
  year    = {2026}
}