FOCUS: DLLMs Know How to Tame Their Compute Bound

AI summary

Overview

Research area: Efficient inference systems for Diffusion Large Language Models (DLLMs), at the intersection of machine learning systems and model architecture analysis.

Technical level: Advanced. The paper assumes familiarity with transformer attention, KV caching, batching, FLOP accounting, and diffusion-based text generation.

Scope: This paper identifies that DLLM decoding is compute-bound due to redundant computation on non-decodable tokens, and proposes FOCUS, an inference system that predicts and evicts those tokens to raise throughput.

What This Paper Is About

Diffusion Large Language Models generate text by denoising a whole block of tokens in parallel, but at each diffusion step only a small fraction of those tokens is actually ready to be decoded. The paper reports that with block size B = 32, the mean proportion of successfully decoded tokens per step is typically only around 10%, meaning roughly 90% of the block-wise computation is redundant. The goal of FOCUS is to identify which tokens are decodable early in the forward pass and skip the rest, so that throughput can keep scaling with batch size instead of saturating on a compute wall.

Key Contributions

1. Identifying the compute-bound nature of DLLM inference. The authors show that DLLMs incur much higher FLOPs per layer than Auto-Regressive LLMs because queries are block-wise, with FLOPs scaling linearly with query size Q (Eq. 1), and that batching gains diminish as compute saturates.

2. Discovering that attention importance predicts decodability. They introduce the "importance delta" — the difference in a token's aggregated attention importance between Layer 1 and Layer 0 — and show empirically that it correlates with whether a token will be decoded.

3. Building FOCUS, a training-free inference system. FOCUS combines Dynamic Budgeting, a Token Eviction strategy guided by the importance delta, and an Intra-Block KV Cache with a Neighbor-Aware Stability Criterion.

4. Demonstrating throughput gains against a production engine. Against LMDeploy, FOCUS achieves up to 3.52× throughput improvement at large batch sizes while preserving or improving generation quality across benchmarks.

Main Findings

• Wasted computation quantified: With block size B = 32, the diffusion process recomputes an entire block to yield only 2.00–4.05 tokens on average, and the mean proportion of decoded tokens is typically around 10%.

• The signal appears at Layer 1, not Layer 0: Importance scores at Layer 0 show negligible difference between decodable and non-decodable tokens, but from Layer 1 onward decodable tokens gain significant attention mass while non-decodable tokens are suppressed.

• The importance delta works as a predictor: Defining ΔI as the attention importance difference between Layer 1 and Layer 0, high-percentile tokens are much more likely to be decoded, while low-percentile tokens mostly remain masked. The authors describe the subtraction as a "Common Mode Rejection" mechanism that filters out non-specific Layer 0 attention.

• Top selection dominates alternatives: When the importance-delta selection is replaced by Random or Bottom-K selection at retention budgets K ∈ {2, 4, 8}, quality follows the hierarchy Top > Random > Bottom across GSM8K, Math500, HumanEval, MBPP and IFEval on both SDAR and LLaDA2.0.

• Quality is preserved or improved: FOCUS matches or outperforms the baseline across confidence thresholds 0.9, 0.8 and 0.7 and expansion factors α ∈ {1.2, 1.5, 1.8}. The default configuration α = 1.5 and Conf = 0.8 yields an average score of 68.37, exceeding the conservative baseline at Conf = 0.9 (67.60).

• Robust at looser thresholds: As the threshold relaxes from 0.9 to 0.7, the baseline drops on Math500 from 64.70 to 54.60, while FOCUS (α = 1.5) maintains 62.20.

• Redundancy ratio reduced sharply: The processed-to-decoded token ratio (N_processed/N_decoded, Layer 2+) falls from 15.02 to 3.12 on SDAR-ShareGPT (79.23% reduction), 14.83 to 3.05 on SDAR-WildChat (79.43%), 7.45 to 2.69 on SDAR-MATH (63.89%), 19.73 to 4.19 on LLaDA2.0-ShareGPT (78.76%), 21.47 to 4.30 on LLaDA2.0-WildChat (79.97%), and 10.13 to 3.04 on LLaDA2.0-MATH (69.99%).

• Throughput scaling restored: The baseline LMDeploy plateaus at roughly 900 tokens/s on SDAR-ShareGPT across batch sizes 32 to 256, whereas FOCUS reaches peak throughputs of 2,272 tokens/s on ShareGPT and 2,324 tokens/s on WildChat at batch size 256 — a 2.32× speedup. On LLaDA2.0 ShareGPT, throughput rises from 2,181 to 3,369 tokens/s at batch size 256.

• Block size matters: FOCUS delivers gains at both B = 16 and B = 64, with the peak speedup of 3.52× at B = 64 where redundancy is heaviest.

• System overhead is negligible: Profiling shows scheduler adaptation and token eviction consume only about 1% of step latency.

• Neighbor-aware caching is necessary: The standard Delayed Cache degrades quality (e.g., 84.92 vs. 89.20 baseline on GSM8K), while the Neighbor-Aware criterion restores stability (88.02), and full FOCUS surpasses the baseline (89.73).

• A caveat on architecture and task: On LLaDA2.0 (MoE with 1.4B active parameters), speedups are more moderate than on dense SDAR, and at batch size 64 a slight regression occurs due to the disablement of multi-loop optimization, where scheduler overhead outweighs FLOPs reduction in long reasoning scenarios.

• Training-free: FOCUS requires no additional training; α > 1 is described as the only hyperparameter it introduces.

Methodology in Plain English

The authors start from the observation that DLLMs spend most of their FLOPs on tokens that will not be decoded in the current step, which makes decoding compute-bound rather than memory-bound. To find a cheap way to identify the tokens worth computing, they borrow an idea from prior attention-sparsity work on Auto-Regressive models — that attention mass concentrates on a few important tokens — but redirect it from compressing the KV cache to evicting query tokens.

They aggregate attention weights across all query tokens in a block and across heads (with a MaxPool1D step following SnapKV) to define a per-token importance score, then subtract the Layer 0 score from the Layer 1 score to get the importance delta. They argue Layer 0 lacks cross-token interaction and mostly reflects priors, while Layer 1 is the earliest point where decodable tokens separate from non-decodable ones.

FOCUS then computes the first two layers for the whole block, ranks tokens by delta, and keeps only a budget K of them. The budget is set dynamically: a historical term based on the recent average number of decoded tokens, scaled by α, combined with a variance-based term counting tokens whose delta is at least one standard deviation above zero. Two structural constraints protect quality — each selected token keeps its immediate predecessor (AR-context preservation), and earlier masked tokens are kept as placeholders so relative positional offsets and KV references stay valid. All later layers operate only on the reduced set of hidden states, with evicted tokens serving as fixed reference KV states. Successfully decoded tokens are frozen in an intra-block KV cache using a delayed mechanism, and the authors add a Neighbor-Aware Stability Criterion that only commits a token's KV states once both it and its right neighbor are decoded. The system is built on top of LMDeploy with custom Triton kernels for irregular memory access, spanning 4,000+ lines. Evaluation uses SDAR-8B-Chat (8B) and LLaDA2.0-mini (16B total, 1.4B active) on a single NVIDIA A100-SXM4-80GB GPU, with default block size 32 and confidence threshold 0.9, sampling 5,000 requests per dataset and a maximum generation length of 2048 tokens.

Why This Matters

Impact on research. The paper reframes DLLM inference as a predictive problem rather than a blind parallel computation problem. It also introduces a new axis of optimization — query-token eviction within a block — that is distinct from the KV cache eviction line of work (H2O, StreamingLLM, SnapKV, Quest in ARLLMs; Sparse-dLLM and d²Cache in DLLMs). It also reports a diagnostic result that Layer 0 attention carries little decodability signal while Layer 1 does, which is relevant to anyone studying early-layer dynamics in bidirectional and diffusion models.

Real-world applications:

• High-throughput serving of DLLM chatbots, where the reported gains are measured on ShareGPT and WildChat style conversational traffic.
• Code generation assistants, since HumanEval and MBPP are among the evaluated quality benchmarks.
• Mathematical reasoning and tutoring tools, evaluated on GSM8K, Math500 and MATH.
• Instruction-following agents, evaluated through IFEval.

Industry relevance. The comparison baseline is LMDeploy, a production-grade engine using continuous batching, PagedAttention and FlashAttention, and the paper reports raw peak throughputs (2,272 and 2,324 tokens/s) and a peak 3.52× speedup at B = 64. The finding that FOCUS is training-free and adds only about 1% of step latency means it can be layered onto an existing serving stack rather than requiring model retraining.

Future Directions

• Better decodability predictors. The authors explicitly call for "more sophisticated decodability predictors" beyond the Layer 0–1 attention delta, shifting from blind redundant computation to proactive predictive decoding.

• Connecting to token editing. The paper notes that LLaDA2.1 proposes token editing as an orthogonal technique, and suggests future work could use early-layer signals to predict which positions require recomputation for either editing or decoding.

• DLLM-native system support. The authors state that DLLM inference still lacks the mature system support of ARLLM stacks, and identify DLLM-native scheduling and GPU-resident state management as important future work, with Appendix F.2 reporting batch-size-1 throughput where scheduling overhead is less amortized.

• Controlled causal validation of quality effects. The authors interpret FOCUS's quality advantages at relaxed thresholds as indirect evidence that importance-based eviction suppresses noisy candidates, and explicitly leave controlled causal validation to future work.

Target Audience

This paper is for systems and machine learning researchers working on efficient LLM inference, particularly those focused on diffusion language models, batching and throughput scaling, or attention-based sparsity. It is also relevant to engineers maintaining production serving stacks who want to understand where DLLM decoding bottlenecks arise and whether token eviction can be added without retraining or quality loss. Readers without prior exposure to transformer attention mechanics, KV caching, or diffusion decoding will find the analysis dense.

Authors’ abstract

Diffusion Large Language Models (DLLMs) offer a compelling alternative to Auto-Regressive models, but their deployment is constrained by high decoding cost. In this work, we identify a key inefficiency in DLLM decoding: while computation is parallelized over token blocks, only a small subset of tokens is decodable at each diffusion step, causing most compute to be wasted on non-decodable tokens. We further observe a strong correlation between attention-derived token importance and token-wise decoding probability. Based on this insight, we propose FOCUS, an inference system designed for DLLMs. By dynamically focusing computation on decodable tokens and evicting non-decodable ones on-the-fly, FOCUS increases the effective batch size, alleviating compute limitations and enabling scalable throughput. Empirical evaluations demonstrate that FOCUS achieves up to 3.52$\times$ throughput improvement over the production-grade engine LMDeploy in large-batch settings, while preserving or improving generation quality across multiple benchmarks.