扩散语言模型的KV Cache量化新方法：TRACE与dLLM-KVQ

为什么扩散语言模型的KV Cache量化更棘手

KV cache量化在自回归语言模型中已有充分研究。在因果注意力机制下，缓存的prefix状态是精确的，量化主要解决的是缓存重建问题。但扩散语言模型（dLLM）打破了这一前提。

Fast-dLLM推理会跨去噪步骤复用KV状态以加速生成，然而双向的token表示会随着被掩码的token逐步揭示而发生变化。这意味着，在被量化之前，复用的缓存可能已经与完全重计算的结果存在偏差。这项研究正是聚焦于这种“缓存复用误差”下的KV cache量化问题。

KV Cache 与自回归解码：在标准自回归（autoregressive, AR）语言模型中，生成是逐 token 串行进行的，每生成一个新 token，就将其对应的 Key/Value 张量追加到缓存中，后续注意力计算直接读取这些缓存而不必重算历史位置。因果注意力保证已缓存的 prefix 状态不会被后续 token 改写，因此缓存一旦写入便与完全重计算的结果完全一致。KV cache 量化的目标就是把这批精确状态压到低比特，同时控制重建误差。Fast-dLLM 等 dLLM 加速方案借用这一思路，把跨去噪步骤“看似不变”的状态缓存下来复用，但双向注意力会让这些状态随掩码 token 的揭示而改变，这正是文章所说“复用缓存可能已偏离重计算结果”的根源。

扩散语言模型与迭代掩码解码：dLLM（如 LLaDA、Dream）不再从左到右生成，而是从全掩码序列出发，经过若干去噪步骤逐步“揭示”被掩码的 token。每一步中模型对整条序列做双向注意力，已揭示 token 的表示会因新揭示的 token 而更新。这与 AR 模型“前缀一旦确定就不再改变”的假设根本不同。因此，跨步骤复用的 KV 状态只是近似值，其误差会与量化误差叠加，这正是 dLLM 的 KV cache 量化需要单独设计的核心原因。

从自回归到扩散：量化假设的迁移失效

自回归模型的量化方法建立在“缓存即精确”的假设上，而dLLM的迭代掩码解码过程让这一假设不再成立。每一步去噪时，部分token被揭示，双向注意力使得已缓存的状态需要根据新揭示的token进行更新。如果直接套用为自回归模型设计的量化方案，复用缓存中的误差会与量化误差叠加，导致输出质量下降。

作者将KIVI和KVQuant适配到Fast-dLLM风格的prefix缓存和dual缓存复用场景，并提出了两种新方法：TRACE和dLLM-KVQ。

KIVI 与 KVQuant：两者都是面向自回归 LLM 的代表性 KV cache 量化方法。KIVI 对 Key 缓存按通道（per-channel）量化、对 Value 缓存按 token（per-token）量化，利用 K 与 V 不同的分布特性；KVQuant 则引入敏感度加权的非均匀量化、离群值处理和校准数据，以减少低比特下的精度损失。它们共同的隐含前提是：用于校准和量化的缓存状态与推理时完全重计算的状态一致。作者把这两类方法迁移到 Fast-dLLM 的 prefix 缓存与 dual 缓存复用场景后，发现复用误差会破坏这一前提，因而需要 TRACE 这类针对揭示轨迹的校准，以及 dLLM-KVQ 这类按层分配比特的策略来弥补。

TRACE：为扩散解码构造校准轨迹

TRACE的核心思路是从prompt-answer对中重建可行的dLLM揭示轨迹（reveal trajectories）。这些轨迹提供了与迭代掩码解码过程相匹配的校准状态。

与自回归模型使用固定前缀校准不同，dLLM的校准需要覆盖token被逐步揭示的动态过程。TRACE通过构造这些轨迹，让量化校准更贴近实际推理时的缓存状态分布，从而减少复用误差带来的影响。

dLLM-KVQ：按层选择高精度K缓存条目

dLLM-KVQ利用TRACE生成的轨迹，根据输出层面和注意力层面的行为指标，选择每层中需要保持高精度的K缓存条目。话说回来，对value缓存应用标准的低比特量化。

这种非对称策略的直觉在于：K缓存对注意力权重的计算更敏感，而V缓存的量化容忍度相对更高。通过仅对部分K条目保留高精度，方法在压缩率和模型质量之间取得平衡。

K/V 缓存的非对称敏感性：在注意力计算中，Key 直接参与 softmax 打分，其微小扰动会改变注意力权重的分布，进而放大到输出；Value 则是在权重确定后做加权求和，误差更容易被平均。已有 AR 量化研究也观察到 K 比 V 更敏感，常对 K 采用更细粒度或更高比特。dLLM-KVQ 把这一直觉与扩散解码的动态性结合：借助 TRACE 轨迹，按输出层和注意力层的行为指标挑出每层最关键的 K 条目保持高精度，其余 K 条目与全部 V 条目走标准低比特量化，从而在压缩率与质量之间取得更好平衡。

实验结果：量化方法中的最高平均分

在LLaDA和Dream模型上，于GSM8K、HumanEval、MATH和MBPP四个基准进行评估，dLLM-KVQ在报告的4-bit和2-bit的prefix/dual缓存设置中，取得了所有量化方法里的最高平均分。

这一结果说明，针对dLLM缓存复用误差设计的量化方案，比直接迁移自回归量化方法更有效。研究也提示，dLLM推理加速的缓存管理需要专门针对双向表示动态性的方法，而非简单的工程适配。

对dLLM推理部署的启示

这项工作的价值在于明确了一个容易被忽视的问题：dLLM的KV cache量化不是自回归量化的直接延伸。缓存复用误差是dLLM特有的误差源，需要在量化校准和比特分配中显式处理。

对于实际部署而言，TRACE提供的轨迹校准思路和dLLM-KVQ的按层混合精度策略，为在有限显存下运行扩散语言模型提供了可参考的方向。随着dLLM在代码生成和数学推理等任务上的应用增多，这类推理优化工作的重要性会进一步显现。