扩散语言模型推理太慢？北大团队：ODB-dLLM破解计算访存双重瓶颈

本研究由北京大学研究团队完成。通讯作者为李萌，北京大学人工智能研究院和集成电路学院助理教授，博导，PKU SEC Lab 负责人，他的研究兴趣集中于高效、安全人工智能加速算法和芯片，旨在通过算法到芯片的跨层次协同设计和优化，为人工智能构建高能效、高可靠、高安全的算力基础。第一作者韦临烨，北京大学集成电路学院博士一年级在读，主要研究方向为多模态高效 AI 系统和加速器设计。

基于扩散的大语言模型 (dLLM) 凭借全局解码和双向注意力机制解锁了原生的并行解码和受控生成的潜力，最近吸引了广泛的关注。例如 Fast-dLLM 的现有推理框架通过分块半自回归解码进一步实现了 dLLM 对 KV cache 的支持，挑战了传统自回归 LLMs 的统治地位。

然而，双向注意力引入的 cache 刷新机制使这一推理过程成为预填充 / 解码交错的模式，将传统自回归模型解码阶段效率的访存瓶颈限制转化为计算 / 访存瓶颈交替限制，占据了推理开销中不可忽视的一部分。此外，认为整个序列 KV 状态在分块内不变的缓存近似方式还一定程度上劣化了模型性能。

针对这一缺陷，来自北大的研究团队提出一种新的 dLLM 推理加速框架 ODB-dLLM（Orchestrating Dual-Boundaries: An Arithmetic Intensity Inspired Acceleration Framework for Diffusion Language Models）。它通过分析现有 dLLM 推理框架中交错的计算和访存瓶颈阶段，引入了自适应长度预测策略和跳跃共享推测解码，以优化 dLLM 在硬件平台上的计算访存特性，最大限度地提高推理效率。

• 论文标题：Orchestrating Dual-Boundaries: An Arithmetic Intensity Inspired Acceleration Framework for Diffusion Language Models
• 论文链接：https://arxiv.org/abs/2511.21759
• Github 链接：https://github.com/PKU-SEC-Lab/ODB-dLLM

研究背景和动机

传统 dLLM 模型中任意两个 token 间的双向注意力破坏了 KV cache 应用的前提，让整个推理过程陷入每一步都执行预填充的计算瓶颈问题，没有取得超越自回归模型的实际推理速度。Fast-dLLM 及类似工作将预设定长度的输出序列划分为多个分块，只在每个分块解码的第一步执行预填充并存储全局的的 KV 状态，随后计算分块内的 tokens 并基于置信度阈值接收。当一个分块内的所有位置都被解掩码，转到下一个分块并刷新 KV cache。然而，这样的推理框架仍然存在以下挑战：

• 交错的计算和访存瓶颈模式都限制了效率。分块半自回归解码的每个新块都会触发全局预填充以刷新 KV cache，占总推理延迟的 30-40%。随着更细粒度的并行解码设计来减少解码步数，计算密集型的预填充阶段将进一步占据主导地位。亟需一个算术强度感知的推理框架联合缓解计算和访存瓶颈。

• 以现有的 dLLM 无法感知实际任务需求，导致预设定输出长度过长。dLLM 中的每个分块都采用双向注意力，以先前和后续令牌为条件。因此，需要固定的响应长度来控制剩余的输出空间。为了确保足够的容量，该长度通常设置为较大的默认值（例如 1024），这会在计算密集型的预填充阶段引入大量开销并导致冗余推理。

• 硬阈值延迟了对潜在正确 token 的接受，导致更长的解码轨迹和更多的步数。在访存密集型的解码阶段，模型并行计算分块内的所有掩码 token，但仅接受置信度超过阈值的部分，而其他 token 则被重新掩码。大量正确 token 由于在当前步的置信度分数低于阈值而被推迟接收，从而导致额外的解码步骤并增加不必要延迟。与此同时，直接降低阈值又会损害模型性能。

ODB-dLLM 框架

ODB-dLLM 推理框架的目标在于协调交错的计算和访存瓶颈限制。为了减轻预设定响应长度带来的大量预填充开销，本文提出了一种自适应长度预测策略来减少冗余计算。此外，本文引入了为 dLLM 特别定制的跳跃共享推测解码方法，可在有限的算术强度下减少访存受限阶段的解码步骤。

自适应长度预测

该策略在每一轮预填充后考察整个响应序列，检测是否存在 [EOS] token。一旦识别出存在这一终止标记，就会提取其置信度分数并将其与预定义的截断阈值进行比较。如果置信度超过阈值，则响应将在该位置被截断。如果未检测到 [EOS] token，或者对应的置信度均低于阈值，则响应长度保持不变。这种渐进式策略最大限度地减少了预填充阶段的算术强度和冗余计算，随着模型的不确定性和多样性逐渐收敛还带来了性能增益。

跳跃共享推测解码

推测解码通过增加计算强度鼓励多样化预测，为访存受限场景下进一步提高解码的并行度提供了解决方案。然而，dLLM 推理的双向注意力将 token 间的因果依赖转化为全局依赖，在传统推测解码中构造的遵循严格顺序的 token tree 被全连通 token graph 所取代。因此，一方面候选 token graph 中任何独立验证的节点对都会解锁其连接边缘，从而实现更灵活的推测解码探索；另一方面任意一个 token 的改变都需要计算一个完整的新块，计算强度随候选预测的增加而迅速增大并演化为访存瓶颈，限制了推测的多样性。基于此，本文提出了两点关键设计：

1. 跳跃接收推测策略。推测解码的草稿块也同时作为中间验证块，实现草稿 token 跨块的跳跃验证和接收，从而在单次解码迭代中解掩码更多 token，在有限的推测块预算下显著提高解码效率。

2. 解码共享推测策略。考虑各推测块间的高度相似性，块内已解码 token 的 KV 状态也可以只在当前块内计算更新，并在多个推测块之间共享。随着解码的进展，推测块的计算成本逐渐降低，为更多推测快的拓展提供了机会。

实验结果：更快、更准

研究团队在两种代表性的开源扩散语言模型LLaDA-Instruct和 LLaDA-1.5 上评估了 ODB-dLLM。在多个数据集上对比 Vanilla model 和 Fast-dLLM 推理框架结果如下：

• 团队测量了预填充长度和延迟的减少，确认了计算密集阶段有意义的算术强度节省。此外，每次预填充迭代期间响应长度的逐步收缩隐式地将模型限制在更准确的答案空间，从而比生成到 [EOS] token 后简单地截断产生更高的精度。

• 针对访存密集阶段的推测解码，实验发现跳跃接收策略增加了每步接受的 token 数量，并通过一次解码中的多段接收减少了并行解掩码 token 间的潜在冲突。解码共享策略的引入进一步减少了解码轮次，同时优先解掩码之前未能接受但置信度较高的 token。跳跃共享推测解码因而在速度和精度上都取得了显著增益。

总结

受预填充和解码阶段异构算术强度的启发，本研究提出了 ODB-dLLM，这是一个协调计算 / 访存双边界以加速 dLLM 推理的框架。在预填充阶段，团队发现预定义的固定响应长度引入了繁重且冗余的计算开销，ODB-dLLM 采用了自适应长度预测机制，可逐步减少预填充开销和不必要的计算。在解码阶段，团队分析了扩散语言模型的计算特性，并提出了一种特别定制的跳跃共享推测解码方法，通过减少解码迭代次数来提高效率。

论文与代码均已公开，欢迎感兴趣的同学阅读、复现以及深入讨论。