模型逐词生成时,会保留前文每层注意力所需的 Key 和 Value。这样下一步无需把整段输入重新计算。上下文越长、并发请求越多,这份 KV 缓存占用的显存就越大。权重已经压到低比特以后,它甚至会成为继续增加批次的主要限制。
我对照了 KIVI 论文的分布分析、量化方法与系统实验。作者没有把 Key 和 Value 当成两块相同的数据。两者参与注意力计算的方式不同,异常值出现的位置也不同,同一种量化维度套在两边会带来很不一样的误差。
Key 的异常值常留在固定通道
KIVI 在 Llama、Falcon 等模型里观察到,Key 缓存有少数通道的数值长期很大。若按每个 Token 单独确定量化范围,一个异常通道会把同一 Token 里其他通道的刻度一起拉宽,小数值能分到的档位随之减少。论文因此按通道给 Key 分组,让异常值的影响留在各自通道内。
论文在 Llama 2 13B 上比较了误差。Key 按 Token 量化时,注意力分数的相对误差接近按通道量化的五倍。这个结果说明低比特位数只交代了存储宽度,分组方向同样会影响模型最后看到的注意力分布。
Value 缓存没有同样稳定的异常通道。它会被注意力权重混合,用某个通道的统一误差尺度处理整段 Token,误差可能跨 Token 影响输出。KIVI 改为按 Token 量化 Value,把误差限制在各个 Token 内。Key 按通道,Value 按 Token,这就是论文所说的不对称量化。
自回归生成还带来一个实现麻烦。Value 每来一个新 Token 就能直接量化后追加。Key 按通道计算尺度时,需要跨多个 Token 才能形成一组。KIVI 把缓存分成已经凑满的量化组和暂时没凑满的残余部分,后者先保留全精度。注意力计算时再把两部分一起使用。这个残余区避免了每生成一个词就重算整组量化参数。
省下显存以后才能扩大批次
KIVI 无需微调模型,把 KV 缓存压到 2 比特。论文在 Llama、Mistral 和 Falcon 上测试生成任务,称质量几乎不变。其 Llama 2 7B 实现把包含模型权重在内的峰值显存缩小为约原来的 1/2.6。省出的空间让批次最多扩大四倍,实测吞吐提高 2.35 倍到 3.47 倍。
这些数字要连着实现条件看。量化和反量化会增加算子工作,硬件是否有合适内核、残余区多大、真实上下文多长,都会改变结果。短输入和低并发下,KV 缓存原本占得不多,压缩它不会自动带来同等速度提升。
部署时可以先画出显存随上下文长度和并发增长的曲线,再分别测试 16 比特、8 比特与更低位宽。除了吞吐,还要看困惑度、任务准确率、长文档问答和输出一致性。显存数字降下来只完成了一半工作,注意力误差有没有落到业务答案里,仍需用真实请求确认。