KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization

LLMs越来越被用于需要大量上下文窗口的文档分析和摘要等应用程序，而在这些大上下文窗口中，KV缓存激活成为推理过程中内存消耗的主要贡献者。量化是一种压缩KV缓存激活的有前途的方法，然而，现有的解决方案无法在超低精度（如小于4位）下准确表示激活。在这项工作中，我们提出了KVQuant，通过包含以下新方法来量化缓存的KV激活来解决这个问题：（i）通道关键量化，我们调整沿着哪个维度量化关键激活以更好地匹配分布；（ii）在旋转位置嵌入之前的关键量化，我们在旋转位置嵌入之前量化关键激活以减轻其对量化的影响；（iii）非均匀KV缓存量化，我们推导出每层灵敏度加权的非均匀数据类型，更好地表示分布；（iv）每向量密集和稀疏量化，我们将异常值单独隔离到每个向量中，以最小化量化范围的偏差；（v）Q-Norm，我们规范化量化质心，以减轻分布偏移，为2位量化提供额外的好处。通过将我们的方法应用于LLaMA、LLaMA-2和Mistral模型，在Wikitext-2和C4上使用3位量化，我们实现了小于0.1的困惑度降低，优于现有方法。我们的方法使得LLaMA-7B模型能够在单个A100-80GB GPU上具有长度高达100万的上下文，并在8-GPU系统上具有长度高达1000万的上下文。

本文旨在解决大文本分析和摘要等应用中，需要大上下文窗口的情况下，KV缓存激活成为推理过程中内存消耗的主要贡献因素的问题。同时，现有的量化方法无法准确地表示超低精度下（如小于4位）的激活情况。

本文提出了KVQuant，通过包括以下几种新方法来量化缓存KV激活，解决了这个问题：（i）通道键量化，调整键激活的维度以更好地匹配分布；（ii）RoPE前键量化，将键激活量化到旋转位置嵌入之前，减轻其对量化的影响；（iii）非均匀KV缓存量化，推导每层灵敏度加权的非均匀数据类型，更好地表示分布；（iv）向量密集-稀疏量化，将异常值单独隔离，以最小化量化范围中的偏差；（v）Q-Norm，规范化量化中心，以减轻分布偏移，为2位量化提供额外的好处。

本文将方法应用于LLaMA，LLaMA-2和Mistral模型，在Wikitext-2和C4上使用3位量化实现了小于0.1的困惑度降低，优于现有方法。该方法使LLaMA-7B模型能够在单个A100-80GB GPU上处理最多100万的上下文长度，并在8个GPU系统上处理最多1000万的上下文长度。

最近的相关研究包括：（1）《Scaling Laws for Neural Language Models》；（2）《On the State of the Art of Evaluation in Neural Language Models》；（3）《The Bottom-up Evolution of Representations in the Transformer: A Study with Machine Translation and Language Modeling Objectives》等。