Paper Source: KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization (arxiv.org)

Background

在LLMs中，随着上下文长度的增加，KV Cache 激活所需的内存成为主要瓶颈，特别是在处理长序列时。这在推理过程中尤为棘手，因为需要高效地存储和处理这些激活。

量化是一种常见的技术，用于通过用更少的位数表示数据来减少内存使用。然而，现有的量化方法在准确表示 KV Cache 的小于4位精度时导致模型性能的显著下降。

KVQuant Approach

研究表明，Key 在某些通道（维度）上容易出现高幅度的异常值，而 Value 在通道和 token 上都有此现象（但不如 Key 表现明显）。

  1. Per-Channel Key Quantization

    核心思想是对每个通道（即每个维度）单独进行量化，而不是对整个 Key 矩阵进行统一的量化。这是因为不同通道中的数据分布可能会有显著差异，尤其是某些通道中的数据可能包含了更重要的信息或更大的数值范围（即所谓的“异常通道”）。通过逐通道量化，可以针对每个通道的具体特性进行更细致的处理，避免信息损失。

    RoPE后，会影响Key矩阵的数据分布，Per-Channel Key Quantization 则在RoPE之前对Key的每个通道进行量化，确保量化能准确反映各通道特征。

  2. Non-Uniform KV Cache Quantization

    根据KV的分布特征，使用不同的量化方式对 KV Cache 进行处理，比如 Key 更适合 Per-Channel 量化，而 Value 更适合 Per-Token 量化。

    • Sensitivity-Weighted Non-Uniform Quantization：对每一层的Key和Value矩阵进行敏感度分析，确定哪些部分对模型性能影响最大。这些部分可能会获得较高的量化精度（位宽更大），而对性能影响较小的部分则会被量化到更低的精度（位宽更小）。
    • Adaptive Layer Quantization Strategy：不同的层可能具有不同的数据特征，基于对每一层的分析，量化方案是动态调整的，而不是一刀切地对所有层使用相同的量化方案。

  3. Per-Vector Dense-and-Sparse Quantization

    将每个向量（例如一个token的Key或Value向量）分解为两部分：稠密部分和稀疏部分。

    • Dense Part：包含向量中大多数数据点，这些数据点通常分布比较均匀，可以通过较低的精度（如3位或4位）进行量化，而不会显著影响模型的性能。
    • Sparse Part：包含向量中的异常值（即偏离正常值范围较大的数据点），这些数据可能对模型的预测非常重要，因此需要用更高的精度来量化，避免丢失关键信息。

    区分每个Vector稀疏和稠密部分的关键是计算出outlier threshold（异常值阈值），计算方式分为两种：Offline Calibration，Online Computation。

  4. Offline Calibration vs. Online Computation