扩散语言模型的 KV 缓存量化

⚡ 一句话总结 AI 生成 · 以原文为准

关键结论

• 将 KIVI 和 KVQuant 适配到 Fast-dLLM 式前缀与双缓存复用,并提出新方法 TRACE 和 dLLM-KVQ。
• TRACE 从提示-回答对重构可行的揭示轨迹,提供与迭代掩码解码匹配的校准状态。
• 在 LLaDA 和 Dream 的 GSM8K、HumanEval、MATH、MBPP 评估中,dLLM-KVQ 在所报告的 4-bit 和 2-bit 前缀/双缓存设置下取得量化方法中的最高平均分。

📄 摘要AI 生成 · 以原文为准中文翻译英文原文

KV 缓存量化在自回归语言模型中已被充分研究,其中缓存的前缀状态在因果注意力下是精确的,量化主要是一个缓存重建问题。扩散语言模型(dLLM)打破了这一假设。快速 dLLM 推理在去噪步骤间复用 KV 状态,但双向 token 表示会随着被遮蔽 token 的揭示而变化,因此在应用量化之前,复用的缓存可能已与完整重计算的结果不同。我们研究了这种缓存复用误差下的 KV 缓存量化。我们将 KIVI 和 KVQuant 适配到 Fast-dLLM 风格的前缀和双缓存复用,并提出了 TRACE 和 dLLM-KVQ。TRACE 从提示-回答对中重构可行的 dLLM 揭示轨迹,提供与迭代掩码解码相匹配的校准状态。dLLM-KVQ 利用这些轨迹,根据输出级和注意力级行为指标,选择需保持高精度的逐层 K 缓存条目,同时对其余部分应用标准低位数值量化。在 LLaDA 和 Dream 上对 GSM8K、HumanEval、MATH 和 MBPP 的评估中,dLLM-KVQ 在所报告的 4-bit 和 2-bit 前缀/双缓存设置下,取得了量化方法中的最高平均分。