通过分块缓存与稀疏选择,提升扩散语言模型长文本推理速度与质量。
Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models

- 将前缀分块缓存,解码时仅选最相关块并利用块内稀疏性
- 在8K-32K上下文中实现9.1-28.0倍加速,长文本任务性能达当前最优
- 用起始标记作周期注意力锚点,解决长序列中丢失中间信息问题
扩散大语言模型(dLLMs)在每个去噪步骤中需重新编码整个前缀,导致计算量随上下文长度呈二次增长,在长文本场景下效率极低。本文提出Prefilling-dLLM,一种无需训练的预填充-解码分离框架:将前缀划分为N个块,一次性缓存其键值(KV)表示,并在解码时基于块内词元稀疏性选取最相关的前K个块。该方法将每步复杂度从全序列长度的二次方降至仅解码长度的二次方,显著降低开销。在LongBench和InfiniteBench上,Prefilling-dLLM在dLLM加速方法中达到领先水平;采用非连续缓存块的并行化注意力核,在8K至32K上下文中实现9.1至28.0倍加速。此外,每块前添加的起始标记作为周期性注意力锚点,有效缓解了‘丢失在中间’现象。代码已开源。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。