​英伟达与港大等合作推出快速 KV 缓存，助力扩散模型提速

　　近日，英伟达、香港大学与麻省理工学院的研究团队联合发布了一种名为 Fast-dLLM 的创新技术，旨在提升扩散语言模型的推理效率。与传统的自回归模型不同，扩散语言模型采用逐步去除文本噪声的方式生成文本，因此能够在一次迭代中生成多个单词，整体效率更高。然而，在实际应用中，许多开源扩散语言模型的推理速度仍然不如自回归模型，主要受限于缺乏键值(KV)缓存支持以及并行解码时生成质量下降。

　　KV 缓存是自回归模型中常用的加速推理技术，通过存储和重用之前计算的注意力状态，显著减少重复计算，从而提高生成速度。但由于扩散语言模型采用双向注意力机制，直接应用 KV 缓存并不容易。Fast-dLLM 架构的创新之处在于将文本生成过程划分为多个块，每个块包含一定数量的 token。通过这种块状生成方式，模型可以在生成一个块之前预先计算并存储其他块的 KV 缓存，进而避免重复计算。

　　尽管 KV 缓存机制有效提升了推理速度，但在并行解码时，生成质量往往会下降。这是因为扩散型模型在解码时假设条件独立性，然而标记之间可能存在复杂的依赖关系。为了解决这一问题，Fast-dLLM 提出了一种基于置信度的平行解码策略。在每个解码步骤中，模型会计算每个标记的置信度，并选择置信度超过阈值的标记进行解码。这种策略确保了高置信度的情况下能够安全进行并行解码，从而维持生成文本的连贯性和准确性。

　　为验证 Fast-dLLM 的性能，研究人员在 NVIDIA A10080GB GPU 上对 LLaDA 和 Dream 两种扩散语言模型进行了全面评估，涵盖了数学推理和代码生成等任务。在 KV 缓存机制的测试中，块大小为32时，模型的吞吐量达到了54.4tokens/s，准确率为78.5%。并行解码测试中，采用动态阈值策略优于固定 token 数基线。整体来看，LLaDA 模型在 GSM8K 任务中仅用 KV Cache 加速3.2倍，并行解码加速2.5倍，两者结合的速度提升达到8.1倍，生成长度达到1024时的端到端加速更是高达27.6倍。所有测试结果显示，Fast-dLLM 在加速的同时，保持了生成质量的稳定。