Abstract

高效部署大型语言模型（LLM）需要将多个请求批量处理以提高吞吐量。随着批量大小、上下文长度或模型尺寸的增加，键值（KV）缓存的大小可以迅速成为GPU内存使用的主要贡献者，也是推理延迟的瓶颈。量化已经成为一种有效的KV缓存压缩技术，但现有方法在非常低的比特宽度下仍然失败。作者观察到，一个键/值激活嵌入的不同通道之间高度相互依赖，多个通道的联合熵的增长速度比它们边缘熵的总和要慢。基于这一洞察，作者提出了耦合量化（CQ），它将多个键/值通道耦合在一起以利用它们之间的相互依赖，并以更信息高效的方式编码激活。大量实验表明，CQ在保持模型质量方面优于或与现有 Baseline 相当。此外，作者证明CQ可以在将KV缓存量化至1位时仍保持模型质量。

1 Introduction

大语言模型（LLM）在各种任务中展示了非凡的泛化能力，包括文本生成、语言翻译和推理，无需特定的微调[25]。这些令人印象深刻的能力使LLM在众多领域找到了应用，如法律[14]、教育[15]和病人护理[35]。然而，LLM的高计算需求和昂贵的部署成本形成了重大障碍，阻碍了它们的广泛应用[14, 3]。特别是，随着LLM向更大模型尺寸[9]和更长上下文长度[34]发展，它们需要更快的有更高内存容量的图形处理单元（GPUs）或其他专用处理器，以便进行有效的推理。因此，开发减少LLM计算成本和内存需求的途径至关重要。

为了加速LLM的推理，关键和值（KV）缓存[20]已被证明是一种有效的技术，且不影响模型质量。在自回归的解码器only LLM中，KV缓存通过在内存中用计算换存储空间：保存当前批次中所有先前 Token 的关键和值激活，以避免在生成下一个 Token 时重新计算它们。然而，由于KV缓存与 Token 数量和批次大小成线性关系，它可以在长上下文或大批次大小设置下迅速超出现有GPU的内存容量。此外，由于过去的键和值激活在序列之间不共享（也许除了一个共同的提示），从GPU内存中读取KV缓存变成了推理的主要瓶颈，而不是下一个 Token 的关注分数和值激活的计算[10]。因此，探索压缩KV缓存的技术有两个主要好处：1. 通过减少KV缓存的内存读取量来加速LLM推理；2. 降低给定批次大小和上下文长度的推理GPU内存需求。现有方法通常通过 Token 驱逐[37, 20]或激活量化[21, 10]来压缩KV缓存。尽管它们可以在中等压缩率（ 压缩或每个浮点数的4位）下保持模型质量，但在高压缩率（ 压缩或每个浮点数的1位）下，模型质量会迅速下降。在这项工作中，作者提出了耦合量化（CQ），一种新颖的KV缓存量化方法，可以在 压缩或每个浮点数的1位下保持模型质量。

作者的方法受到一个观察的启发：同一键/值激活嵌入中的不同通道高度相互依赖。因此，编码多个通道的键/值激活嵌入比独立量化每个通道更为信息高效。现有的KV缓存量化方法采用每个通道或每个 Token 的量化策略，这些策略未能利用通道之间的相互依赖，并在高压缩率下遭受模型质量的灾难性下降。作者提出的方法通过联合量化多个通道来利用通道间的相互依赖，并在大多数情况下比现有方法更好地保持模型质量，尤其是在低比特宽度设置下。作者的贡献总结如下：

1. 作者实证观察到，LLM中的同一键/值激活嵌入的不同通道之间共享了大量的依赖或互信息，这是现有KV缓存压缩方法没有利用的关键洞察。
2. 作者提出了耦合量化（CQ），一种利用联合编码多个通道降低熵的新颖KV缓存量化方法。
3. 通过广泛的实验，作者证明了CQ相对于最具有竞争力的现有方法的有效性。此外，作者展示了CQ在极端KV缓存压缩水平1位时保持模型质量的能力。

LLM Attention and KV Cache

仅解码器的基于Transformer的大规模语言模型采用 Mask 自注意力机制[32]，其中当前 Token 的激活仅依赖于之前的 Token ，而不会受到未来 Token 的影响。这一特性使得可以并行训练下一个 Token 预测目标，并产生了KV缓存技术以实现高效的推理解码。考虑在LLM的单个注意力头中解码第个 Token 的步骤。第个 Token 的输入嵌入（一个列向量），，经过三个不同的变换成为键、 Query 和值激活嵌入，其中变换由线性投影和位置编码（如RoPE[29]）组成。第个 Token 的注意力输出嵌入计算为：

因此，计算当前 Token 的输出嵌入需要所有先前 Token 的键和值激活嵌入，和，其中。这些嵌入在先前的解码步骤中被缓存于内存中，作为_KV缓存_，以避免重复计算并减少推理延迟。KV缓存的大小可以计算为个浮点数，其中是批大小，是每个序列中的 Token 数，是模型中的层数，2代表键和值，是键/值注意力的头数，是单个键/值激活嵌入头中的通道数。随着批大小、上下文长度或模型尺寸的增加，KV的大小很快就会超出有限的GPU内存。

The von Neumann Bottleneck of KV Cache

方程1中的注意力计算主要受到GPU内存带宽的限制，这被称为冯·诺伊曼瓶颈，因为全局内存访问与计算的比例较低。GPU的运算速度远快于从全局内存读取，先前的研究已经表明，通过 Kernel 融合[5]避免不必要的全局内存读写可以加速注意力计算。在LLM注意力的解码阶段，计算当前 Token 的输出需要从全局内存中获取所有前序 Token 的缓存键和值嵌入，这导致了算术操作与全局内存读取的比例较低[10]。此外，每个批次中的每个序列都保留自己的KV缓存，这导致GPU的并行处理能力利用率不高。因此，KV缓存压缩算法可以缓解冯·诺伊曼瓶颈并提高LLM推理效率，即使这种方法在解压缩或去量化时引入了额外的计算开销。由于在注意力中GPU计算核心大多数时间都因为KV缓存内存访问而停滞，量化方法可以有效地减少内存读取，同时由于去量化计算的延迟可以忽略不计。

Channel-wise Quantization

现有的KV缓存量化方法[10, 21]采用通道量化方法对键（keys）进行处理，而对值（values）采用标记（token）量化方法，这种做法是基于观察到的某些键通道在幅度上表现出异常模式，而值通道则没有。通道量化和标记量化相似，不同之处在于学习量化质心的方向。在非均匀通道量化中，为每个通道学习一组质心。假设是一个键或值激活矩阵，而表示的第个通道。那么，非均匀-比特通道量化旨在独立地为的每个通道学习一组质心，通过以下目标

其中将中的每个值量化到中最近的质心。

3 Methodology

在本节中，作者使用信息论来激发作者的 Proposal ，并介绍了用于KV缓存压缩的耦合量化（CQ）方法。

Motivations

作者的方法受到了信息论中概念[28]的启发。作者将关键/值激活嵌入中的每个通道视为一个随机变量。通道中的信息量（或不确定性）可以通过熵来衡量，定义为，其中是概率密度函数，是的支持集。衡量了无损编码通道所需的理论比特数，因此可以用来判断一个通道的“可量化”程度：如果$H(X_{1})<h(x_{2})$，那么通道$x_{1}$可能比通道$x_{2}$量化为更少的比特，同时达到相同的量化误差。< p=""></h(x_{2})$，那么通道$x_{1}$可能比通道$x_{2}$量化为更少的比特，同时达到相同的量化误差。<>

作者的洞察是，来自同一关键/值激活嵌入的不同通道可能是相互依赖的，这比独立编码多个通道时所需的比特数要少。两个通道中的总信息量（或不确定性）通过联合熵来衡量，定义为，其中是联合概率密度函数。两个通道的联合熵是它们的边际熵之和与它们的互信息之差，即，其中是用于衡量两个随机变量相互依赖性的非负量。因此，作者有

这意味着联合编码两个通道所需的比特数不超过独立编码它们所需的总比特数。以前的工作已经证明深度神经网络[11]和基于注意力的网络[7]倾向于产生低秩嵌入，这表明LLM中关键/值嵌入的通道可能表现出高度的相互依赖性。

因此，测量多通道的联合熵与它们边际熵之和在关键值激活嵌入中的差异是有益的。如果这种差异显著，那将表明一起编码这些通道比独立编码它们在信息效率上更高。然而，由于不知道它们的概率密度函数，所以无法推导出通道的确切熵或联合熵。因此，作者采用“分箱”技巧[17]来估计熵。作者首先通过在数据集上保存KV缓存来观察关键值通道的实证分布，并将每个通道的支持分割成大小相等的箱子。然后，将每个通道的值离散化到它们所属箱子的索引。最后，个通道的联合熵用黎曼和估计，

其中是分箱或离散化的支持，是经验概率密度函数。具体来说，作者将LLaMA-7b的关键值嵌入通道划分为大小为的不重叠组，其中，并估计每组联合熵和边际熵之和。每个通道的支持被划分为16个大小相等的箱子。图1显示了在WikiText-2数据集的262k个 Token 上，LLaMA-7b三层平均每组估计的联合熵和边际熵之和的平均值和标准差。作者只展示最大组大小为4，因为增加组大小需要指数级地保存更多关键值嵌入，以避免空箱子并保持估计质量。如图1所示，边际熵之和以线性速率增长，而联合熵以次线性速率增长得更慢。这意味着随着联合量化通道数量的增加，编码所需的总信息量减少。这一现象是作者提出方法的动机基础。

除了展示估计的边缘熵和联合熵之外，作者还展示了LLaMA-7b模型在WikiText-2数据集上，键（key）和值（value）激活嵌入通道之间的皮尔逊相关系数。相关系数捕捉两个随机变量之间的线性依赖关系。图2展示了8个层次的相关矩阵的热力图，而所有层次的相关矩阵则在附录的第六节中呈现。键和值通道表现出高度的线性依赖性，正如相关系数的高幅度所示，它们显然不是独立分布的。

Coupled Quantization

受到不同键/值通道间高度依赖性的启发，作者提出了耦合量化（Coupled Quantization，CQ），这是一种针对压缩LLM KV缓存的信息高效量化方法。与现有的KV缓存量化方法不同，它们是按通道或按标记进行量化，CQ对键和值执行通道耦合量化。更具体地说，键或值激活嵌入的通道被划分为大小相等、互不重叠的连续通道组。每组中的通道是“耦合”的，因为它们共同量化并共享一个量化代码。对于每组耦合通道，学习一组独特的多通道质心，其中每个质心的维数等于该组中的通道数。在量化键或值激活嵌入时，每组耦合通道被量化到与其最近的质心（按L2距离计算）。作者使用CQ-

在CQ中，通过利用统一聚类或基于二阶信息聚类的方法，在校准数据集上为每个通道组学习多通道质心。具体来说，对于CQ-_c_c_b_b配置的统一质心学习，通过以下目标独立地为每个通道组学习一组质心：

其中是包含第组所有耦合通道的的子矩阵，将每个列向量量化到质心中最近的质心，以L2距离为标准。作者使用k-means算法[22]和k-means++初始化[1]来优化目标。

大型语言模型（LLMs）对某些权重的量化精度比其他权重更为敏感[16]。为了更好地保持模型质量，CQ的质心应该学习偏向于保留更重要激活的精度。为此，作者利用Hessian的近似来进行基于二阶信息的质心学习。更具体地说，作者使用Fisher信息矩阵的对角线来识别更有影响力的关键/值激活，并指导质心学习过程。这种方法由Li等人[18]提出，并由Kim等人[16]用于通道权重量化，作者将它扩展到多通道CQ。为了执行Fisher引导的质心学习，作者首先在一个校准数据集上保存一个关键/值激活矩阵及其梯度，其中是训练损失函数。作者使用梯度矩阵的元素平方来近似Hessian矩阵，即。作者使用Fisher信息矩阵对角线元素之和作为每组通道耦合激活重要性的度量，并使用以下目标为第个通道组获得质心集：

作者利用加权k-means来优化目标。质心学习过程和质心存储的开销在4.3节中讨论。

作者通过将LLaMA-7b KV缓存压缩到1位和2位来验证作者提出的通道耦合和Fisher引导的质心学习的有效性，并在图4中展示了WikiText-2下不同CQ配置的困惑度结果和量化误差（在层上平均的）。实验设置在4节给出。随着耦合通道数量的增加，困惑度和量化误差显著改善，接近FP16 Baseline 性能。尽管Fisher引导的质心学习增加了量化误差，但它更好地保留了显著的激活并实现了更低的困惑度。

4 Experiments

在本节中，作者进行了大量实验以验证作者提出的CQ方法在KV缓存压缩中的有效性。作者首先介绍实验设置，包括硬件、软件、评价指标、数据集和所用到的 Baseline 。然后，作者展示详细的实证结果并提供讨论。最后，作者进行消融研究以验证作者 Proposal 中每个组件的有效性。

硬件测试平台 实验在运行Ubuntu 20.04的Linux服务器上进行，配备了2个AMD EPYC 7742 CPU、1.5TB内存和4个NVIDIA A100 40GB GPU。

软件实现 作者基于PyTorch [24] 和 HuggingFace Transformers库 [33] 实现了CQ的软件部分。

评价指标和基准 作者在采用不同KV缓存量化算法的情况下，评估了5个流行的开源LLM在各类基准上的质量。考虑的5个LLM分别是：1. LLaMA-7b, 2. LLaMA-13b [30], 3. LLaMA-2-7b, 4. LLaMA-2-13b [31], 5. Mistral-7b [13]。作者使用困惑度指标评估LLM在WikiText-2 [23] 和 C4 [26] 数据集上的质量，并在零样本设置下使用准确度评估3个基准：WinoGrande [27], PIQA [2], 和 ARC Challenge [4]。困惑度是在LLM的最大上下文长度（LLaMA为2048，LLaMA-2为4096，Mistral为8192）下，在数据集的测试集上进行评估。

作者将作者提出的方法与未压缩的FP16 KV缓存以及具有竞争力的KV缓存量化方法进行了比较，包括：1. 均匀整数（INT）量化（不分组且组大小为128），2. 正常浮点（NF）量化 [6]（不分组且组大小为128），3. KVQuant [10]（没有稀疏异常值，并将1%的异常值以稀疏格式存储）。KVQuant-b-1% 是一种密集和稀疏方法，除了在推理过程中需要密集矩阵乘法外，还需要额外的稀疏矩阵乘法，这引入了额外的计算开销。KVQuant和作者提出的CQ都需要在校准数据集上学习质心，对于这两种方法，作者都使用了相同的校准集，该集合包含16个序列（每个序列有2048个标记）的WikiText-2。校准在WikiText-2的训练集上执行一次，而困惑度和准确度是在不同数据集和基准的测试集上评估的。对于每种方法，作者报告了每个浮点数（FPN）的位数以衡量压缩率，该压缩率是通过将每个标记的量化KV缓存的位数除以每个标记的未压缩KV缓存中FPN的数量得出的，不包括质心、缩放因子和零点的常数存储开销。

Results

表2：在不同KV缓存量化方法下，不同位宽的C4数据集上的困惑度。INT、NF和KVQuant（不包括1b和1b-1%）的结果来自[10]。作者的方法CQ一致地优于非密集与稀疏量化方法，并且在与密集与稀疏方法KVQuant-b-1%比较时表现更佳或相当。

WikiText-2上的困惑度结果在表1中呈现，C4的结果在表2中呈现。CQ一致地优于非密集与稀疏量化方法，特别是在低位宽区域。尽管使用了更低的位宽，CQ的表现仍然与密集与稀疏量化方法KVQuant-b-1%相当或更好。密集与稀疏量化方法由于对激活异常值的额外稀疏矩阵乘法而引入了额外的推理开销，这在GPU上效率不高，而CQ没有这个限制。

不同基准测试上的准确度结果展示在表3中。CQ一致地优于位宽为1和2时的非密集与稀疏 Baseline KVQuant-b，并且在与密集与稀疏 Baseline KVQuant-b-1%比较时表现更佳或相当。

作者进行了一系列消融实验来研究作者提出方法中每个组成部分的有效性。消融实验的结果展示在表4中。作者在2位每FPN的情况下，使用CQ评估了Mistral-7b和LLaMA-2-13b这两个模型在WikiText-2上的困惑度，并通过比较不同通道耦合数量下的统一质心学习和Fisher引导的质心学习。Fisher引导的质心显著提高了模型质量，这一点通过更低的困惑度得到了证明。无论是统一质心还是Fisher引导的质心，随着耦合通道数量的增加，困惑度都有所改善。因此，作者提出的通道耦合和Fisher引导的质心学习方法对于保持模型质量都是有效的。

Overhead of Centroid Learning and Storage

在本节中，作者讨论了CQ的质心学习计算开销和质心存储的内存开销。CQ的质心学习过程包含许多独立的k-means运行，这在CPU上可能会非常耗时。因此，作者利用GPU实现来加速学习过程。在所有的实验中，作者使用k-means++初始化，并在单个GPU上运行100次k-means迭代以获得质心。存储质心的内存开销可以计算为 个16位浮点数，其中 是层数，2用于键和值， 是键/值注意力头的数量， 是单个头键/值激活嵌入的通道数， 是量化代码的位宽。不同CQ配置和模型的详细学习与内存开销在表5中给出。CQ可以轻松扩展到大型模型规模，且具有较低的学习和内存开销。

5 Related Works

大型语言模型（LLMs）的高内存需求和计算要求对有效推理提出了巨大挑战。在训练后的模型权重量化已被证明可以有效减少推理延迟和内存需求。"GPTQ [8]" 将基于近似海森的权重量化扩展到大规模模型，而 "AWQ [19]" 保持了LLM中的显著权重，以实现更好的量化模型质量。"SqueezeLLM [16]" 利用基于敏感度的非均匀聚类和密集-稀疏量化来保留显著权重和异常值。除了权重量化之外，KV缓存压缩方法也有效于提高推理效率。"Scissorhands [20]" 和 "H2O [37]" 通过凸显不重要的标记，只存储关键标记，实现了KV缓存的压缩同时保持了模型质量。"KIVI [21]" 对关键激活逐通道量化，对值激活逐标记量化，并使用残差来获得更好的模型质量。"KVQuant [10]" 提出了基于敏感度的量化和密集-稀疏量化用于KV缓存。"FlashAttention [5]" 通过融合 Kernel 消除不必要的全局内存读写，提高了基于注意力模型在GPU上的推理效率，而 "NoMAD-Attention [36]" 通过利用寄存器内Shuffle加速CPU上的注意力计算。"Product Quantization [12]" 是一种近似最近邻搜索方法，通过将向量空间分解为低维子空间的笛卡尔积，并在每个子空间内共同量化维度来实现向量的压缩。

6 Conclusion

在本工作中，作者提出了耦合量化（CQ）方法，通过压缩KV缓存来使LLM推理更加高效，这在长上下文或大批量设置中是延迟和吞吐量的瓶颈。作者观察到，不同的键/值激活嵌入通道之间存在高度依赖性，这尚未被现有的压缩方法利用。受这一洞察的启发，作者提出了通道耦合方法，以利用通道间的依赖性，从而实现键/值激活的信息效率更高的编码。此外，作者还提出了由Fisher引导的质心学习方法，以更好地保留显著激活和模型质量。大量实验表明，在相同的量化位宽下，作者的方法在模型质量方面通常优于现有方法。此外，CQ能够在将KV缓存量化到1位时，合理地保持模型质量。

图7：在WikiText-2上，4个LLaMA-7b层的键激活嵌入中的通道对二维散点图。

图8：在WikiText-2上，4个LLaMA-7b层的值激活嵌入中的通道对二维散点图。

图8：在WikiText-2上，4个LLaMA-7b层的值激活嵌入中的通道对二维散点图。