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

📄 arXiv: 2606.10537v1 📥 PDF

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

分类: cs.CL

发布日期: 2026-06-09

备注: Technical Report

🔗 代码/项目: GITHUB

💡 一句话要点

提出Prefilling-dLLM以解决长上下文推理中的计算复杂性问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散大语言模型 长上下文推理 稀疏注意力 解码加速 自然语言处理

📋 核心要点

  1. 现有的扩散大语言模型在处理长上下文时，计算复杂度随上下文长度呈平方级增长，导致效率低下。
  2. 本文提出Prefilling-dLLM，通过将前缀分块并缓存其KV表示，利用稀疏选择机制来优化解码过程。
  3. 在LongBench和InfiniteBench上，Prefilling-dLLM展示了优于现有方法的性能，显著提升了解码速度和质量。

📝 摘要（中文）

扩散大语言模型（dLLMs）在每个去噪步骤中重新编码整个前缀，导致计算量随着上下文长度呈平方级增长，尤其在长上下文场景中变得不可行。为此，本文提出了Prefilling-dLLM，这是一种无训练的预填充-解码分解框架，将前缀分为N个块，缓存其KV表示，并选择与解码相关性最高的前K个块进行解码。结果表明，稀疏预填充在减少每步复杂度的同时，能够超越密集注意力。Prefilling-dLLM在LongBench和InfiniteBench上实现了dLLM加速方法中的最优质量，并且在8K到32K上下文中，通过并行解码非连续缓存块的注意力内核实现了9.1到28.0倍的加速。

🔬 方法详解

问题定义：本文旨在解决扩散大语言模型在长上下文推理中计算复杂度过高的问题。现有方法在每个去噪步骤中需重新编码整个前缀，导致计算量随上下文长度平方增长，尤其在长上下文场景中变得不可行。

核心思路：提出Prefilling-dLLM框架，通过将前缀分为多个块，缓存其KV表示，并在解码时选择最相关的块进行处理，从而减少计算复杂度。这样的设计旨在提高长上下文推理的效率，同时保持模型的性能。

技术框架：Prefilling-dLLM的整体架构包括前缀分块、KV表示缓存和稀疏选择机制。首先，将输入前缀分为N个块，并在训练前缓存每个块的KV表示；然后，在解码阶段，根据相关性选择前K个块进行解码。

关键创新：最重要的创新在于引入了稀疏预填充机制，显著降低了每步解码的复杂度，从而使得解码过程的复杂度仅与解码长度的平方成正比，而非整个序列长度。

关键设计：在实现中，设计了高效的注意力内核，能够并行处理非连续缓存块的KV表示。此外，序列开始的标记被添加到每个块中，作为周期性注意力锚点，避免了在长序列中出现的“迷失在中间”现象。具体的参数设置和损失函数设计在实验部分进行了详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在实验中，Prefilling-dLLM在LongBench和InfiniteBench上达到了最优的质量，并且通过并行解码非连续缓存块的注意力内核，实现了9.1到28.0倍的速度提升，显著优于现有的dLLM加速方法。

🎯 应用场景

Prefilling-dLLM可广泛应用于自然语言处理领域，尤其是在需要处理长文本的任务中，如文档摘要、长篇对话生成等。其高效的解码机制和优越的性能将为实际应用提供更快的响应时间和更高的质量，推动相关技术的进一步发展。

📄 摘要（原文）

Diffusion large language models (dLLMs) re-encode the entire prefix at every denoising step, causing recomputation that scales quadratically with context length and becomes prohibitive for long-context scenarios. We propose Prefilling-dLLM, a training-free prefill-decode disaggregation framework for dLLMs that partitions the prefix into N chunks, caches their KV representations once, and selects the top-K most relevant chunks with intra-chunk token sparsity for decoding, showing that sparse prefilling can outperform dense attention while reducing per-step complexity from quadratic in the full sequence length to quadratic only in the decode length. On LongBench and InfiniteBench, Prefilling-dLLM achieves state-of-the-art quality among dLLM acceleration methods, and an attention kernel that parallelizes decoding over the non-contiguously cached chunk KV yields 9.1--28.0x speedup at 8K--32K contexts. We further show that beginning-of-sequence tokens prepended to each chunk act as periodic attention anchors that eliminate the lost-in-the-middle phenomenon. Code is available at https://github.com/menik1126/Prefilling-dLLM.