💬 NLP

Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models

The paper introduces Prefilling-dLLM, a training-free framework that accelerates long-context inference in diffusion language models by caching and sparsely selecting relevant prefix chunks to reduce computational complexity from quadratic in full sequence length to quadratic in decode length, thereby achieving state-of-the-art speedups and mitigating the lost-in-the-middle phenomenon.

Original authors: Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

Published 2026-06-10

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a massive jigsaw puzzle, but instead of looking at the whole picture at once, you have to re-examine the entire box of pieces every single time you place just one new piece on the table.

That is essentially how Diffusion Large Language Models (dLLMs) currently work when handling long stories or documents. Every time the model tries to generate a new word, it re-reads and re-processes the entire history of the conversation from the very beginning. As the story gets longer, this "re-reading" becomes so slow and expensive that it's like trying to run a marathon while carrying a heavy backpack that gets heavier with every step.

The paper introduces a new method called Prefilling-dLLM to fix this. Here is how it works, using some everyday analogies:

1. The Problem: The "Re-Reading" Trap

In traditional AI models, once you read a page of a book, you remember it and only focus on the new page you are reading. But in these Diffusion models, the computer forgets the book and re-reads the entire book from page 1 to page 1,000 every time it writes a single new word.

The Result: If the book is short, it's fine. If the book is 32,000 pages long, the computer spends 99% of its time just re-reading the old pages and only 1% actually writing new ones.

2. The Solution: The "Smart Librarian" System

The authors propose a system called Prefilling-dLLM that acts like a super-efficient librarian. Instead of re-reading the whole library every time, the librarian does two things:

Step A: The "Prefill" (Organizing the Library)

Before the writing starts, the librarian takes the massive input text (the "prefix") and chops it up into chunks (like chapters in a book).

The Trick: The librarian reads each chapter once, creates a "summary card" (called a KV cache) for it, and puts it on a shelf.
The Innovation: The librarian doesn't read every single word in every chapter. They use a special trick to only keep the most important sentences from each chapter, throwing away the fluff. This makes the summary cards much smaller and faster to handle.

Step B: The "Decode" (Writing the Story)

Now, when the model needs to write the next word, it doesn't re-read the whole library.

The Selection: The model asks, "Which chapters are actually relevant to what I'm writing right now?" It uses a smart scoring system to pick only the top few chapters (the most relevant chunks) from the shelf.
The Efficiency: It ignores the other 90% of the library that isn't needed for this specific sentence. It only looks at the relevant "summary cards" it made earlier.

3. Why This is a Game-Changer

The paper claims this approach is a massive speedup because:

No More Re-Reading: The heavy lifting of reading the long text happens once at the beginning.
Smart Skipping: During the actual writing phase, the model only looks at a tiny fraction of the text (the top relevant chunks), rather than the whole thing.
The Speed: On long contexts (8,000 to 32,000 words), this method is 9 to 28 times faster than previous methods, while still getting the answers right.

4. Solving the "Lost in the Middle" Mystery

There is a known problem with long AI models called "Lost in the Middle." Imagine a person reading a long list of facts; they remember the first few and the last few perfectly, but they completely forget the stuff in the middle.

The Paper's Fix: The authors found that by chopping the text into chunks and adding a special "anchor" token (like a chapter heading) to the start of each chunk, the model can find information anywhere in the text, even the middle. It's like having a table of contents that works perfectly, so the model never gets lost, no matter how long the story is.

5. The "Chunk Size" Balance

The paper also tested how big these "chapters" should be.

Too Big: If the chunks are huge, the model might miss important details in the middle of the text.
Too Small: If the chunks are tiny, the computer spends too much time managing the list of chunks.
The Sweet Spot: They found that for very long texts, using many smaller chunks works best to keep the accuracy high.

Summary

Prefilling-dLLM is like switching from a student who re-reads a 500-page textbook every time they answer a quiz question, to a student who creates a smart, condensed study guide once, and then only flips to the specific pages they need when answering.

The result? The AI can handle much longer conversations and documents without slowing down to a crawl, and it doesn't forget the important details hidden in the middle of the text.

Technical Summary: Prefilling-dLLM

Problem Statement

Diffusion Large Language Models (dLLMs) offer a promising alternative to autoregressive (AR) models by enabling parallel token generation through iterative denoising. However, dLLMs face a critical inefficiency in long-context scenarios: at every denoising step, the model must re-encode the entire input prefix to compute self-attention. This results in a computational complexity that scales quadratically with the total sequence length ( $O((L_p + L_d)^2 \cdot T)$ , where $L_p$ is prefix length, $L_d$ is decode length, and $T$ is denoising steps). Unlike AR models, where the prefill phase is compute-bound and the decode phase is memory-bound, dLLM inference remains compute-bound throughout all iterations because the full sequence (prefix + response) is jointly processed at every step. This repeated recomputation of the prefix becomes prohibitive for long-context applications (e.g., 8K–32K+ tokens).

Methodology: Prefilling-dLLM

The authors propose PREFILLING-DLLM, a training-free, prefill-decode disaggregation framework designed to eliminate prefix recomputation and introduce sparsity. The method operates in three main stages:

1. Chunked Prefill and KV Caching

Instead of processing the entire prefix as a single block, the input prefix is partitioned into $N$ non-overlapping chunks of fixed size $C$ .

Independent Processing: Each chunk is processed independently with intra-chunk bidirectional attention. A special Beginning-of-Sequence (BOS) token is prepended to each chunk to act as a delimiter.
Predictive Scoring: To determine which chunks are relevant without ground-truth targets, the framework generates pseudo-labels ( $m$ $m$ ) via a preliminary diffusion pass. It then scores each chunk using two signals:
- Self-Information Score ( $s_I$ ): The negative log-likelihood of the query window conditioned on the chunk.
- Pseudo-Label Score ( $s_P$ ): The likelihood of the pseudo-labels conditioned on the chunk and query.
Selection: The top- $K$ most relevant chunks are selected based on the combined score ( $s = s_I + s_P$ ). The KV representations of these selected chunks are cached.

2. Two-Level Sparsity

The framework introduces sparsity at both the chunk and token levels to reduce complexity:

Inter-chunk Sparsity: Only the top- $K$ chunks participate in the attention mechanism during decoding, rather than the full set of $N$ chunks.
Intra-chunk Sparsity: Within each selected chunk, tokens are ranked by their cumulative bidirectional attention weight with the query tokens. Only the top- $B$ most relevant tokens per chunk are retained in the KV cache, evicting irrelevant tokens.

3. Decoding with Static KV Reuse

During the iterative denoising process:

The cached KV representations of the selected $K$ chunks remain static and are reused across all $T$ denoising steps.
The model only recomputes the KV representations for the query and the evolving response tokens.
This reduces the per-step complexity from $O((L_p + L_d)^2 \cdot T)$ to $O(N \cdot C^2 + (L_d^2 + K \cdot C) \cdot T)$ , effectively decoupling decoding cost from the full prefix length.

Attention Kernel Optimization

To address memory bottlenecks caused by loading non-contiguous chunk KV caches, the authors implement a Split-S FlexAttention kernel. This kernel parallelizes decoding over the non-contiguously cached chunk KVs by partitioning the KV dimension and processing chunks independently before merging results via log-sum-exp reduction, avoiding the need to gather data into contiguous memory.

Key Findings and Results

Performance and Efficiency

Speedup: On LongBench and InfiniteBench, PREFILLING-DLLM achieves 9.1× to 28.0× speedup at context lengths of 8K to 32K compared to standard Transformers baselines.
Quality: The method achieves state-of-the-art quality among dLLM acceleration methods. On LongBench, it achieves an average score of 34.59 (Dream-7B) and 37.02 (UltraLLaDA), outperforming baselines like Fast-dLLM, Sparse-dLLM, and dKV-Cache.
Long-Context Scaling: On InfiniteBench (contexts >128K), the method achieves 43.62% average accuracy on Dream-7B, significantly surpassing the strongest baseline (Fast-dLLM v2 at 30.32%). Notably, it achieves 95.42% accuracy on the Passkey retrieval task and 70.00% on Number retrieval.

Addressing "Lost-in-the-Middle"

Phenomenon: The paper observes that dLLMs exhibit a "lost-in-the-middle" phenomenon under context extrapolation, where accuracy degrades for information located in the middle of the context, showing a monotonic decay bias toward the response rather than the U-shaped curve seen in AR models.
Mitigation: By selecting chunks based on predictive relevance rather than positional proximity, PREFILLING-DLLM eliminates this bias. The method maintains high exact-match accuracy across all positions in the context, even at 32K and 128K lengths.

Mechanism Insights

Attention Anchors: The BOS tokens prepended to each chunk act as periodic attention anchors. Analysis shows these tokens distribute attention mass uniformly across the context (forming stable, low-magnitude ridges) rather than degenerating into "attention sinks" that absorb disproportionate attention mass. This mechanism enables position-invariant retrieval.
Chunk Size Trade-off: There is a trade-off between chunk size and retrieval accuracy. Smaller chunks (e.g., 256 tokens) maintain high accuracy even at 128K contexts, while larger chunks (e.g., 4096 tokens) degrade sharply beyond 32K. Multi-chunk selection is essential for long-context inference.

Significance and Claims

The paper claims that PREFILLING-dLLM bridges the gap between the compute-bound nature of dLLM inference and the efficiency of prefill-decode disaggregation architectures used in AR models. Its primary contributions are:

Training-Free Disaggregation: It introduces a framework that caches prefix KV representations once and reuses them, eliminating the quadratic recomputation cost of the prefix at every denoising step without requiring model retraining.
Sparse Prefilling: It demonstrates that sparse prefilling (selecting relevant chunks and tokens) can outperform dense attention in dLLMs, reducing complexity while improving or maintaining generation quality.
System-Level Optimization: It provides a specialized attention kernel (Split-S FlexAttention) that leverages the chunked structure of the prefix to achieve significant latency reductions (up to 10.2× over vanilla FlexAttention) during the decode phase.
Robustness to Long Contexts: It shows that this approach mitigates the "lost-in-the-middle" phenomenon and enables dLLMs to scale effectively to 128K+ tokens with high retrieval accuracy, a capability previously unexplored in dLLM acceleration.

The authors note limitations, including the static nature of chunk selection (no dynamic re-selection during decoding), the need for task-specific tuning of chunk size and $K$ , and the current evaluation being limited to English benchmarks on specific model architectures (Dream-7B and UltraLLaDA).