KVQuant

KVQuant 项目介绍

KVQuant 是一种高效的密钥-值（KV）缓存量化方法，致力于实现长上下文长度语言模型（LLM）的推理。其核心在于通过创新的低精度量化技术，突破了长上下文长度推理中的内存瓶颈。

项目背景

在处理长上下文长度推理时，大模型会产生巨大的KV缓存需求，这对计算资源提出了极高的要求。为了解决这个问题，KVQuant 引入了多种策略以提升低精度KV缓存量化的准确性。这些策略包括：

每通道预旋转位置编码（Pre-RoPE）密钥量化，以更好地适配密钥中的异常通道。
非均匀量化（NUQ），用于更好地表示非均匀激活。
密集与稀疏量化结合，以缓解数值异常值对量化难度的影响。

这些创新使得 KVQuant 能够在单个 A100-80GB GPU 上处理 LLaMA-7B 模型的 1M 上下文长度，如果使用 8-GPU 系统甚至可以处理 10M 以上的上下文长度。这一突破显著提高了大型语言模型在实际应用中的可行性。

大型世界模型推理

LWM 是一种最近开发的训练方法，支持最大 1M 上下文长度的模型。然而，由于KV缓存的庞大，该模型的推理成本极为高昂。KVQuant 的出现，使得这一推理过程在单个 A100 GPU 上变得更加高效。

附加方法改进

KVQuant 进一步提升了支持长上下文长度推理的能力，具体改进包括：

GPU 并行 topK 支持：通过额外的内核拓展了开源系统，可以进行多输入符号的并行打包。
密钥异常值限制：在每个符号中的异常值数量限制下，进行校准和推理，从而设计出更高效的内核。
注意力汇机制量化：根据注意力汇论文的观察，模型倾向于将注意力集中在第一个符号上，且在校准和推理时，保持一部分初始密钥和值使用 fp16 格式。

安装与代码结构

KVQuant 的代码库包含五个子文件夹，每个文件夹都附有安装所需环境的 README 文件：

gradients：用于计算费舍尔信息，以支持新模型的量化。
quant：用于模拟量化和评估实验。
deployment：用于通过压缩向量进行高效推理。
lwm：用于运行和评估量化的 LWM 模型。
benchmarking：用于内核基准测试。

未来计划

KVQuant 将继续优化内核代码，并计划合并更新后的推理部署代码，以及增加长上下文长度评估的脚本。

引用

如果您在工作中使用了 KVQuant 的代码库，敬请引用下述论文：

@article{hooper2024kvquant,
  title={KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization},
  author={Hooper, Coleman and Kim, Sehoon and Mohammadzadeh, Hiva and Mahoney, Michael W and Shao, Yakun Sophia and Keutzer, Kurt and Gholami, Amir},
  journal={arXiv preprint arXiv:2401.18079},
  year={2024}
}