HCAttention：通过异构注意力计算实现大规模语言模型（LLMs）的极端键值（KV）缓存压缩

《Neurocomputing》：HCAttention: Extreme KV cache compression via heterogeneous attention computing for LLMs

　　

杨东泉|杨一帆|于晓天|齐先标|肖荣

中国深圳Intellifusion有限公司

摘要

使用大型语言模型处理长上下文输入是一个重大挑战，因为在推理过程中需要存储大量的键值（KV）缓存。因此，随着序列长度的不断增加，高效的KV缓存管理策略对于平衡性能和资源利用变得至关重要。然而，当内存减少超过85%时，现有方法的性能会明显下降，这是在压缩比和信息损失之间难以权衡的问题。我们提出了HCAttention，这是一种异构注意力计算框架，它可以动态选择KV缓存而无需永久性移除。我们的方法结合了GPU上的键量化、全局近似注意力评分以及CPU驻留值的延迟检索，以减少GPU KV缓存的内存使用量，同时保持模型质量。该方法与现有的Transformer架构兼容，且不需要对模型进行微调。在LongBench基准测试上的实验结果表明，我们的方法在将KV缓存内存占用量减少到原始大小的25%的同时，仍保持了全注意力模型的准确性。值得注意的是，仅使用12.5%的缓存，它在极端GPU KV缓存预算下仍能保持强大的性能。

引言

大型语言模型（LLMs）[1]、[2]、[3]、[4]、[5]、[6]、[7]、[8]在广泛的自然语言处理任务中展示了先进的性能。它们的核心自回归机制逐个生成文本标记，每个新标记都依赖于之前的所有标记。随着现代应用（如多轮对话[9]、[10]、文档理解[11]、[12]和AI代理[13]、[14]）对上下文窗口的需求扩展到数十万甚至数百万个标记，支持这样的极端序列长度暴露出了严重的效率瓶颈。

这种低效率的主要原因是自注意力机制[15]及其相关的键值（KV）缓存。虽然自注意力的计算成本随序列长度N的增加而呈二次方增长，但KV缓存的内存占用量却呈线性增长（$$' role="presentation">$$），因为它必须为上下文中的每个标记存储一对键值向量。这种内存消耗很快变得难以承受。例如，对于像OPT-175B[16]这样的模型，处理128个长度为2048的序列批次仅KV缓存就需要大约950GB的内存，几乎是模型参数本身所需内存的三倍。因此，KV缓存成为了主要的内存瓶颈，限制了批量大小和可实现的上下文长度，从而阻碍了其在实际场景中的部署。

尽管之前的工作主要集中在改进注意力计算本身——通过算法或位置编码调整[17]、[18]、[19]、[20]等方法——但这些技术对解决KV缓存的主要内存成本作用有限。因此，对于极长上下文的推理，我们的方法直接将重点放在KV缓存管理上，研究了三种互补策略：KV缓存选择、量化和卸载。

• •
  KV缓存选择[21]、[22]仅保留“重要”的KV对，而永久丢弃其他KV对。然而，如果被丢弃的标记对后续生成至关重要，这种方法可能会导致不可逆的准确性损失。
• •
  KV缓存量化[23]、[24]通过使用低精度（例如8位或4位）来减少内存。然而，其压缩比本质上是有限的（通常为50%-25%），并且仍然可能引入准确性下降。
• •
  KV缓存卸载[25]将缓存从高带宽的GPU内存转移到较低级别的存储（如CPU RAM）。虽然有效减少了GPU内存压力，但引入了显著的数据传输开销。

总之，每种方法都在压缩比、模型准确性和通信开销之间呈现了一个基本的权衡。当内存减少超过85%时，极低位的量化和KV缓存选择会导致明显的性能下降，这是在压缩比和信息损失之间难以权衡的问题。卸载则会产生巨大的通信开销。因此，开发一种能够在不牺牲模型质量或推理效率的情况下大幅缩小KV缓存的技术仍然是一个关键且未解决的挑战。

我们的方法基于一个关键见解，该见解来自之前工作的观察[22]、[26]：LLMs对KV缓存中性能下降的敏感性在不同注意力层和头中差异显著。这意味着静态的、一刀切的压缩策略是不理想的。相反，一种动态且可逆的KV缓存管理策略，避免永久性移除KV对，对于保持模型性能至关重要。基于此，我们认为单一种技术是不够的。我们需要一个统一的框架来全面协调这些策略，动态适应不同层的需要，以最大化压缩效果而不牺牲准确性。

基于这些原则，我们提出了HCAttention，这是一种用于高效长上下文推理的异构注意力框架。HCAttention基于两个核心策略：

动态且可逆的KV缓存移除：一种动态的、非永久性的移除策略，通过累积幅度标准进一步从GPU缓存中剔除最不重要的KV对。这确保了最重要的信息在GPU的不同层和头中仍然易于访问，而较不重要的数据则被暂时移动，而不会永久丢失。

不对称的K-V管理策略：该框架根据键和值的作用不同地处理它们。对于注意力评分至关重要的键进行量化，并保留在GPU上进行快速计算。而仅在最终聚合步骤中需要的值则被卸载到CPU内存中。

•

我们提出了一种动态且可逆的KV缓存移除算法。与那些可能造成永久信息损失的方法不同，我们的策略智能地管理GPU缓存以保持模型保真度。

•

我们开发了一种不对称的K-V管理流程，将GPU上的量化键评分与CPU驻留值的延迟检索相结合，减少了GPU内存使用量，同时保持了全局评分阶段。

•

我们在长上下文基准测试上提供了全面的实验，表明HCAttention将KV缓存占用量减少到原始大小的12.5%，同时保持了与标准全注意力模型相当的性能和准确性。

部分摘录

相关工作

在以下三个类别中探索了可扩展LLM部署的高效KV缓存管理：KV缓存选择、KV缓存量化和LLM推理的异构计算。

HCAttention

HCAttention引入了一种异构注意力框架，其新颖之处在于协调了三个组件：（1）用于全局近似注意力评分的GPU上量化的键表示，（2）用于最终聚合的CPU驻留值的延迟检索，以及（3）避免永久性移除令牌的动态可逆移除策略（图1）。

实验结果

在本节中，我们使用Llama模型在长上下文基准测试上评估了HCAttention。我们还进行了消融研究，以分析每个组件的单独贡献。

结论

在本文中，我们介绍了HCAttention，这是一种旨在解决长上下文LLM推理中关键KV缓存内存瓶颈的异构框架。我们的方法基于两个核心原则：在量化键上进行的全局近似注意力评分，随后是CPU驻留值的延迟检索，以及一种动态可逆的移除算法，该算法管理GPU缓存而不会造成永久性信息损失。我们的理论分析证实，这种设计显著减少了

CRediT作者贡献声明

杨东泉：软件、方法论、形式分析、数据整理、概念化。杨一帆：写作——审阅与编辑、撰写——初稿、可视化、验证、方法论、形式分析。于晓天：写作——审阅与编辑、项目管理、方法论、概念化。齐先标：写作——审阅与编辑、验证。肖荣：监督、项目管理。

利益冲突声明

作者声明他们没有已知的竞争性财务利益或个人关系可能会影响本文报告的工作。

杨东泉是深圳Intellifusion科技有限公司的算法研究员。他毕业于中山大学数据科学与计算机科学学院，主修统计机器学习。他的研究兴趣包括自然语言处理、大型模型训练、大型模型量化和剪枝，专注于在边缘-云协作场景中高效部署大型模型。