在处理长上下文的大型语言模型(LLM)推理任务时,KV缓存(Key-Value Cache)已成为主要的成本之一。在自回归解码过程中,缓存的大小随上下文长度、批处理大小和模型深度线性增长。当批处理量大、上下文长度达到十万级别Token,并有数十个并发请求时,KV缓存会占用GPU内存的绝大部分。因此,压缩KV缓存是直接增加批处理量和减少内存流量的有效方法。
最直接的解决方案是量化。然而,将KV缓存推至2比特(INT2)精度一直以来都极具挑战性。现有方法要么会导致模型精度大幅下降,要么需要与现有分页KV缓存系统不兼容的自定义服务布局。Together AI提出的OSCAR(Offline Spectral Covariance-Aware Rotation)系统,成功解决了这两个难题。
为什么INT2 KV缓存量化如此困难?
KV激活中包含通道维度上的异常值。少数通道的数值异常大,而大多数通道的行为则相对正常。当应用只有四个可表示级别的INT2量化时,这些异常值会主导量化比例因子。这导致量化器将其大部分范围浪费在稀疏的峰值上,而正常数值则被压缩到仅一两个有效级别中,从而严重降低注意力机制的质量。
基于旋转的量化方法通过应用固定的正交变换(通常是Hadamard变换)来重新分配异常值能量到所有通道,从而解决这一问题。这种方法在INT4量化上效果尚可。但在INT2精度下,一个更深层次的问题依然存在:旋转变换是数据无关的。它虽然可以平滑激活范围,但并不知道注意力机制实际读取的方向。均匀地分散量化误差与将误差推向低重要性方向是不同的。在只有四个量化级别的INT2情境下,这种区别决定了模型是否能正常工作。
OSCAR有何不同之处?
OSCAR的关键观察点在于,量化前应用的旋转应该源自注意力统计信息本身,而非仅仅基于KV激活的原始分布。
对于键(Keys),真正重要的下游误差不是K的欧几里得重构误差,而是注意力逻辑值(attention logits)的误差。研究团队证明,此误差为:‖QK⊤ − QK̂⊤‖²F = tr((K − K̂)Q⊤Q(K − K̂)⊤)。这里的权重矩阵是查询协方差Q⊤Q,而不是K⊤K。在查询具有高能量的方向上,量化误差在逻辑值中会被放大。OSCAR通过校准集估计经验查询协方差CQ = (1/N) Σ qn⊤qn,对其进行特征分解,并使用特征向量UQ作为键的旋转基。
对于值(Values),相关误差体现在注意力输出SV中。这取决于注意力分数矩阵S如何加权每一行值。研究团队定义了分数加权值协方差CS = (1/N) V⊤S⊤SV。经过S聚合后仍然较大的方向,是量化误差传播通过的方向。OSCAR使用CS的特征向量US作为值的旋转基。
最终的复合旋转...