Prefilling-dLLM：扩散语言模型中长上下文推理的预测性预填充

# Prefilling-dLLM：扩散语言模型长上下文推理中的预测性预填充  
来源：https://arxiv.org/html/2606.10537  
Jing Xiong¹, Qi Han¹, Shansan Gong¹, Yunta Hsieh², Chengyue Wu¹, Chaofan Tao¹, Chenyang Zhao³, Ngai Wong¹  
¹香港大学，²密歇根大学安娜堡分校，³LMSYS Org  

###### 摘要  
扩散大语言模型（dLLMs）在每个去噪步骤中都会重新编码整个前缀，导致计算量随上下文长度呈二次增长，在长上下文场景中变得不可行。我们提出Prefilling-dLLM，一个针对dLLMs的**免训练**预填充-解码分离框架。该框架将前缀划分为N个固定大小的块，仅缓存其KV表示一次，并在解码时通过块内令牌稀疏性选取最相关的K个块。实验表明，稀疏预填充的性能可媲美密集注意力，同时将每步复杂度从全序列长度的二次方降低至仅解码长度的二次方。在LongBench和InfiniteBench上，Prefilling-dLLM在dLLM加速方法中达到了最先进的质量。通过一个注意力内核实现非连续缓存块KV上的并行解码，在8K–32K上下文中实现了9.1–28.0倍的加速。我们进一步发现，预置在每个块开头的起始序列令牌（BOS tokens）充当周期性的注意力锚点，消除了“中间迷失”现象。¹¹我们的代码已开源：https://github.com/menik1126/Prefilling-dLLM。  

## 1 引言  
扩散大语言模型（dLLMs）已成为自回归（AR）模型的一种有前景的替代方案，通过迭代去噪能够并行生成多个令牌（Nie等人，2025 (https://arxiv.org/html/2606.10537#bib.bib25)；Ye等人，2025a (https://arxiv.org/html/2606.10537#bib.bib24)；Sahoo等人，2024 (https://arxiv.org/html/2606.10537#bib.bib56)；Austin等人，2021 (https://arxiv.org/html/2606.10537#bib.bib52)）。与AR模型从左到右逐个生成令牌不同，dLLMs同时破坏和重建整个序列，支持灵活的生序方式并有望实现更快的推理（Wu等人，2025b (https://arxiv.org/html/2606.10537#bib.bib22)；Wang等人，2025a (https://arxiv.org/html/2606.10537#bib.bib44)）。然而，这种范式在长上下文场景中引入了严重的低效问题：整个输入前缀必须参与每个去噪步骤，尽管其表示在迭代过程中基本保持不变。  

在自回归LLM服务中，**预填充-解码分离**架构（Zhong等人，2024 (https://arxiv.org/html/2606.10537#bib.bib1)）将预填充和解码阶段分配给不同的GPU集群，利用它们不同的计算特征（预填充是计算密集型，而解码是内存密集型），从而最大化硬件利用率和服务吞吐量。相比之下，dLLM推理本质上始终是计算密集型：由于整个序列（前缀+解码）必须在每个去噪步骤中联合处理，每次迭代都会对整个组合序列执行一次完整的前向传播，使得工作负载主要由矩阵乘法主导，而非**内存带宽**。这种计算密集型特性在所有去噪迭代中持续存在，与仅每步附加一个令牌的AR解码不同。  

近期关于dLLM加速的工作探索了KV缓存策略（Ma等人，2026 (https://arxiv.org/html/2606.10537#bib.bib5)；Liu等人，2025b (https://arxiv.org/html/2606.10537#bib.bib23)；Nguyen-Tri等人，2025 (https://arxiv.org/html/2606.10537#bib.bib30)）和稀疏注意力机制（Wang等人，2025b (https://arxiv.org/html/2606.10537#bib.bib43)；Song等人，2025 (https://arxiv.org/html/2606.10537#bib.bib28)；Jiang等人，2025 (https://arxiv.org/html/2606.10537#bib.bib31)），但尚无工作探索通过分离预填充和解码阶段来避免跨去噪迭代重复进行长上下文计算。  

我们的关键洞察是：在长上下文dLLM推理中，输入前缀在每个去噪迭代中被冗余处理，然而响应令牌对前缀的注意力表现出强烈的局部性偏差，且这种偏差随步骤增加而加剧——只有一小部分前缀令牌被积极关注。基于这一观察，我们提出Prefilling-dLLM（Diffusion LLMs的预填充），它在专用的预填充阶段计算前缀KV缓存一次，并在所有解码步骤中重复使用，无需重新计算。具体而言，我们将前缀划分为N个固定大小的块（块大小为C），并采用块内注意力，将预填充复杂度从O(L_p²)降低到O(N·C²)，并支持跨设备并行处理。在解码阶段，我们通过检索增强生成（Jiang等人，2024 (https://arxiv.org/html/2606.10537#bib.bib41)；Lai等人，2025 (https://arxiv.org/html/2606.10537#bib.bib42)；Xu等人，2025 (https://arxiv.org/html/2606.10537#bib.bib48)；Yuan等人，2025 (https://arxiv.org/html/2606.10537#bib.bib49)）仅保留一小部分相关块，将复杂度从O((L_p+L_d)²·T)降低到O(N·C² + (L_d² + K·C)·T)，其中K为所选块的数量，T为去噪步数。  

我们在LongBench和InfiniteBench上评估Prefilling-dLLM，在8K–32K上下文中实现了9.1–28.0倍的加速，同时在dLLM加速方法中达到了最先进的质量。我们的贡献如下：  
- **我们提出了Prefilling-dLLM**，一个针对dLLMs的**免训练预填充-解码分离**框架。通过一次性预填充前缀KV缓存并在所有去噪迭代中共享，消除了重复计算，并实现了随上下文长度显著增长的加速。  
- **我们引入了稀疏预填充**，选择相关的块和令牌，将复杂度从O((L_p+L_d)²·T)降低到O(N·C² + (L_d² + K·C)·T)。结合优化的注意力内核，在缓存的块KV上实现解码并行化，在32K上下文中获得了高达28倍的端到端加速。  
- **我们证明**，预置在每个块开头的BOS令牌充当周期性注意力锚点，在不引入注意力汇的情况下缓解了dLLM中的“中间迷失”现象。  

## 2 相关工作  

### 2.1 扩散语言模型  
扩散模型已通过各种形式从连续域扩展到离散文本生成。早期工作探索了基于词嵌入的连续扩散（Li等人，2022 (https://arxiv.org/html/2606.10537#bib.bib55)；Gong等人，2022 (https://arxiv.org/html/2606.10537#bib.bib54)）和基于离散令牌的掩码扩散（Austin等人，2021 (https://arxiv.org/html/2606.10537#bib.bib52)；He等人，2023 (https://arxiv.org/html/2606.10537#bib.bib53)；Sahoo等人，2024 (https://arxiv.org/html/2606.10537#bib.bib56)）。最近，掩码离散扩散已扩展到大规模语言模型（Gong等人，2025 (https://arxiv.org/html/2606.10537#bib.bib18)）：LLaDA（Nie等人，2025 (https://arxiv.org/html/2606.10537#bib.bib25)）展示了掩码扩散在8B参数规模下可匹配自回归模型，而Dream（Ye等人，2025a (https://arxiv.org/html/2606.10537#bib.bib24)）和MDLM（Sahoo等人，2024 (https://arxiv.org/html/2606.10537#bib.bib56)）进一步验证了该范式的有效性。后续工作聚焦于扩展规模（Bie等人，2025 (https://arxiv.org/html/2606.10537#bib.bib17)；Gong等人，2025 (https://arxiv.org/html/2606.10537#bib.bib18)）、偏好对齐（Zhu等人，2025 (https://arxiv.org/html/2606.10537#bib.bib50)），以及将dLLM扩展到长上下文（Liu等人，2025a (https://arxiv.org/html/2606.10537#bib.bib34)；He等人，2025 (https://arxiv.org/html/2606.10537#bib.bib35)）和多模态设置（You等人，2025 (https://arxiv.org/html/2606.10537#bib.bib14)）。尽管取得了这些进展，dLLM在长上下文场景中的效率问题仍未充分探索。  

### 2.2 dLLM的高效推理  
在自回归LLM中，稀疏注意力方法如MInference（Jiang等人，2024 (https://arxiv.org/html/2606.10537#bib.bib41)）、DCA（An等人，2024 (https://arxiv.org/html/2606.10537#bib.bib73)）、FlexPrefill（Lai等人，2025 (https://arxiv.org/html/2606.10537#bib.bib42)）、XAttention（Xu等人，2025 (https://arxiv.org/html/2606.10537#bib.bib48)）和NSA（Yuan等人，2025 (https://arxiv.org/html/2606.10537#bib.bib49)）通过自适应或块稀疏模式降低长上下文注意力成本，而StreamingLLM（Xiao等人，2024 (https://arxiv.org/html/2606.10537#bib.bib10)）、H2O（Zhang等人，2023 (https://arxiv.org/html/2606.10537#bib.bib37)）和SnapKV（Li等人，2024 (https://arxiv.org/html/2606.10537#bib.bib38)）通过仅保留重要条目来压缩KV缓存。然而，这些技术针对的是从左到右生成过程中自然构建KV缓存的因果注意力，并不直接适用于dLLM中的双向注意力——其中不存在此类缓存。  

针对dLLM，Fast-dLLM（Wu等人，2025b (https://arxiv.org/html/2606.10537#bib.bib22)）和Fast-dLLM v2（Wu等人，2025a (https://arxiv.org/html/2606.10537#bib.bib51)）通过重用前次迭代的键值表示，在去噪步骤间引入了KV缓存。dKV-Cache（Ma等人，2026 (https://arxiv.org/html/2606.10537#bib.bib5)）提出了自适应缓存，根据令牌置信度选择性更新KV条目。SparseD（Wang等人，2025b (https://arxiv.org/html/2606.10537#bib.bib43)）、Sparse-dLLM（Song等人，2025 (https://arxiv.org/html/2606.10537#bib.bib28)）、d2Cache（Jiang等人，2025 (https://arxiv.org/html/2606.10537#bib.bib31)）、Focus-dLLM（Long等人，2026 (https://arxiv.org/html/2606.10537#bib.bib72)）和LoSAX等人（2026 (https://arxiv.org/html/2606.10537#bib.bib71)）利用固有的注意力稀疏性进行动态缓存驱逐。然而，所有这些方法都在标准推理循环内运行，每步仍然处理整个序列。我们的工作则从系统层面将前缀计算与迭代解码分离，并在静态前缀KV缓存上应用稀疏块检索。  

### 2.3 预填充-解码分离  
在自回归LLM服务中，预填充是计算密集型，而解码是内存密集型。DistServe（Zhong等人，2024 (https://arxiv.org/html/2606.10537#bib.bib1)）利用这种不对称性，将两个阶段分配给不同的GPU集群。Mooncake（Qin等人，2024 (https://arxiv.org/html/2606.10537#bib.bib2)）通过分布式缓存池在预填充和解码节点之间传输KV缓存，SPAD（Zhang等人，2025 (https://arxiv.org/html/2606.10537#bib.bib3)）为每个阶段设计了专门的硬件，Semi-PD（Hong等人，2025 (https://arxiv.org/html/2606.10537#bib.bib4)）引入了一种结合分离计算和统一存储的混合方法。这一原理尚未应用于dLLM，因为dLLM的每个去噪步骤都对整个序列执行完整前向传播，推理始终是计算密集型的。我们的工作通过一次性计算前缀KV缓存并在所有去噪迭代中重复使用来填补这一空白，并进一步分析了缓存引入的潜在内存瓶颈。  

## 3 预备知识：掩码扩散模型  
掩码扩散语言模型（dLLMs）定义了一个前向噪声过程（Sahoo等人，2024 (https://arxiv.org/html/2606.10537#bib.bib56)；Gong等人，2025 (https://arxiv.org/html/2606.10537#bib.bib18)；Ye等人，2025a (https://arxiv.org/html/2606.10537#bib.bib24)），通过将令牌替换为特殊的[MASK]令牌来逐步破坏离散令牌序列 x₀ = (x₁, ..., x_L)。在每个扩散时间步 t ∈ [0,1]，每个令牌以概率 t 独立地被掩码，得到噪声序列 xₜ。逆向（去噪）过程由一个神经网络 p_θ(x₀|xₜ) 参数化，该网络根据部分掩码的输入预测原始的干净令牌。训练时，模型优化以最小化掩码位置上的交叉熵损失：  
L = E_{t, x₀, xₜ} [ - Σ_{i: xₜⁱ=[M]} log p_θ(x₀ⁱ | xₜ) ].  (1)  

推理时，模型从一个完全掩码的序列开始，在 T 个去噪步骤中逐步揭开令牌。在每一步，模型同时预测所有掩码位置，并根据调度策略揭开一部分高置信度的预测。这种并行解码使 dLLM 每步能生成多个令牌，但每步模型需要在整个序列（前缀+响应）上执行完整的自注意力，导致每次迭代的计算成本与总长度成正比。  

## 4 动机  

### 4.1 dLLM 中的“中间迷失”现象  
自回归LLM存在“中间迷失”现象（Liu等人，2024 (https://arxiv.org/html/2606.10537#bib.bib65)），即对位于上下文中间的信息检索准确率下降。我们使用一个位置可控的多文档QA任务评估dLLM是否具有相同的偏差，并发现三个关键观察：（i）在原生训练范围（256–2K令牌）和YaRN×2外推（4K）下，Dream-7B在所有位置都达到完美准确率；（ii）进一步外推（8K、16K、32K）引入了新出现的位置敏感性（图1 (https://arxiv.org/html/2606.10537#S4.F1)），准确率偏向于更靠近响应的位置，这与AR LLM中两端（开头和结尾）均占优的U型曲线不同；（iii）在dLLM中，双向注意力产生单调衰减：靠近响应的令牌无论其绝对位置如何都获得强烈关注，而远处的令牌被均匀忽视。这种由局部性驱动的退化激发了我们基于块的选择性检索策略。  

参见图注  
图1：Dream-7B的“中间迷失”评估（训练长度=2K）。通过YaRN缩放进行上下文外推。原生范围（256–2K）和YaRN×2（4K）在所有位置实现EM=1.0。YaRN×4（8K）、YaRN×8（16K）和YaRN×16（32K）显示逐渐增加的退化。每个位置使用30个样本评估；每上下文长度均匀设置10个位置。  

### 4.2 注意力衰减的局部性  
参见图注  
图2：响应令牌到前缀令牌的注意力权重随距离（令牌间隔数）的衰减，在不同去噪步骤测量。注意力随距离快速衰减，表现出强局部性偏差。在后续去噪步骤中，随着令牌预测趋于稳定，这种衰减更加明显。  

我们进一步分析了Dream-7B在去噪过程中的注意力模式，以理解响应令牌如何关注前缀。我们测量了响应令牌对前缀令牌的平均注意力权重作为距离（分隔它们的令牌数）的函数。我们观察到三个关键发现（图2 (https://arxiv.org/html/2606.10537#S4.F2)）：（i）注意力权重随距离迅速衰减，响应令牌将其大部分注意力集中到附近的前缀令牌上；（ii）这种衰减