论文
Prefilling-dLLM:扩散语言模型中长上下文推理的预测预填充
Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models
摘要
扩散 大语言模型 (dLLM) 在每个去噪步骤中重新编码整个前缀,导致重新计算与上下文长度成二次方缩放,并且对于长上下文场景来说变得令人望而却步。我们提出了 Prefilling-dLLM,一种用于 dLLM 的 无需训练 预填充解码分解框架,它将前缀划分为 N 个块,将其 KV 表示缓存一次,并选择具有块内词元稀疏性的前 K 个最相关的块进行解码,这表明稀疏预填充可以优于密集注意力,同时将每步复杂度从完整序列长度的二次降低到仅在解码长度中的二次。在 LongBench 和 InfiniteBench 上,Prefilling-dLLM 在 dLLM 加速方法中实现了最先进的质量,并且对非连续缓存块 KV 进行并行解码的注意力内核在 8K--32K 上下文中产生 9.1--28.0 倍的加速。我们进一步表明,每个块前面的序列开头标记充当周期性注意力锚,消除了中间丢失的现象。代码可在 https://github.com/menik1126/Prefilling-dLLM 获取。