CommVQ：用于KV缓存压缩的可交换矢量量化

大语言模型（LLMs）在需要长上下文的应用中得到越来越广泛的应用，但随着上下文的增长，键值（KV）缓存经常成为GPU上的内存瓶颈。为了解决这个问题，我们提出了交换向量量化（CommVQ），以显著减少长上下文LLM推理的内存使用。我们首先引入了带有轻量级编码器和码本的加性量化，用于压缩KV缓存，该缓存可以通过简单的矩阵乘法进行解码。为了进一步降低解码时的计算成本，我们设计了与旋转位置嵌入（RoPE）可交换的码本，并使用期望最大化（EM）算法进行训练。这使得解码能够高效地集成到自注意力机制中。我们的方法通过加性量化实现了高精度，并通过RoPE可交换码本实现了低开销。在长上下文基准测试和GSM8K上的实验表明，我们的方法通过2位量化将FP16 KV缓存大小减少了87.5%，同时优于最先进的KV缓存量化方法。值得注意的是，它支持1位KV缓存量化，且精度损失极小，使得LLaMA-3.1 8B模型能够在单张RTX 4090 GPU上以128K的上下文长度运行。源代码可在以下网址获取： https://github.com/UMass-Embodied-AGI/CommVQ。