Unifying KV Cache Compression for LargeLanguage Models with LeanKV——使用LeanKV统一大型语言模型的KV缓存压缩

这篇文章介绍了一种名为LeanKV的统一键值（KV）缓存压缩框架，旨在提高大型语言模型（LLMs）的服务效率，同时保持模型准确性。文章的主要内容包括：

1. 问题背景：LLMs在推理过程中需要存储大量的键值（KV）缓存，这导致了高内存消耗和服务成本。现有的KV缓存压缩方法（如量化和剪枝）存在局限性，无法有效处理键和值向量的不同重要性以及动态的注意力模式。

2. LeanKV的创新：

   • 异质KV量化：键向量以更高精度存储，值向量以较低精度存储，以反映键向量对注意力计算的更大影响。

   • 每头动态稀疏性：根据每个注意力头和每个请求的令牌重要性动态分配内存，适应不同头和请求之间的动态注意力模式。

   • 统一KV压缩：结合混合精度量化和选择性剪枝，实现模型准确性和内存效率之间的平滑权衡。

3. 系统优化：

   • 统一分页：将内存划分为固定大小的页面，支持不同精度的令牌存储，最大化内存带宽利用率。

   • GPU上的并行内存管理：通过并行化内存管理操作，减少内存管理的开销。

4. 实验结果：LeanKV在多个模型和基准测试上进行了评估，结果显示在不损失准确性的情况下，KV缓存压缩了3.0倍到5.0倍，并在允许5%准确性损失的情况下压缩高达11.0倍，吞吐量提高了1.9倍到6.9倍。

5. 结论：LeanKV通过其创新的压缩技术和系统优化，显著提高了LLM的服务效率，为大规模语言模型的推理提供了高效的解决方案。

LeanKV通过异质量化、动态稀疏性和统一压缩策略，有效减少了KV缓存的内存占用，提升了LLM的推理效率，同时保持了模型的准确性。这里是自己的论文阅读记录，感兴趣的话可以参考一下，如果需要阅读原文的话可以看这里，如下所示：

摘要：
大型语言模型（LLMs）表现出卓越的性能，但由于键值（KV）缓存的高内存需求，导致服务成本高昂。现有的KV缓存压缩方法，包括量化和剪枝，存在一些局限性，例如对键和值的统一处理以及跨注意力头的静态内存分配。为了解决这些挑战，我们提出了LeanKV，一个统一的KV缓存压缩框架，通过以下三项创新来提高LLM的服务效率而不影响准确性：（1）异质KV量化，以更高的精度存储键，以反映它们在注意力计算中的更大影响；（2）每头动态稀疏性，根据每个头和每个请求的令牌重要性分配内存；（3）统一KV压缩，集成混合精度量化和选择性剪枝，以实现模型准确性和内存效率之间的平滑权衡。为了有效支持这些技术，LeanKV引入了系统优化，包括统一分页和GPU上的并行内存管理。在vLLM上实现后，LeanKV将KV缓存压缩了3.0倍到5.0倍而不损失准确性，并在准确性损失小于5%的情况下压缩高达11.0倍，吞吐量提高了1.9倍到2.5倍，最高可达6.9倍。

1. 引言

大型语言模型（LLMs）如GPT和Gemini已经展示了显著的应用潜力，涵盖了聊天机器人、编程、数学和医疗辅助等领域。尽管它们表现出色，但由于模型规模庞大，托管LLMs的成本高昂，需要大量的硬件资源。随着LLMs的广泛应用，提高服务效率变得至关重要。

LLMs基于Transformer架构，该架构使用注意力机制根据输入提示和先前生成的令牌顺序生成令牌。在注意力机制中，每个令牌被转换为查询、键和值特征向量。注意力分数通过查询和键向量的点积计算得出，使模型能够权衡不同令牌的重要性。随后，输出基于注意力分数计算为值向量的加权和。为了捕捉令牌之间的多样化关系，注意力机制通常分为多个头，每个头独立计算注意力，并将它们的输出连接和转换以创建更丰富的输入数据表示。

为了避免冗余计算，LLM推理框架通常将中间键和值张量缓存在内存中，通常称为KV缓存。KV缓存的大小随序列长度和并发请求数量线性增长，通常占内存消耗的90%以上。随着最先进的模型支持更长的序列，KV缓存已成为LLM服务效率的关键瓶颈，限制了内存中可以容纳的并发请求数量，并由于内存带宽限制增加了注意力计算的延迟。

认识到KV缓存在限制服务效率中的关键作用，研究人员已经研究了各种压缩技术，主要集中在量化和剪枝上。量化通过将浮点特征值转换为整数来减少KV缓存的大小，而剪枝则通过消除被认为不重要的令牌来减少令牌序列长度。尽管这些技术有潜力，但它们在实际应用中存在效率和实用性的局限性。首先，现有的量化方法对所有令牌的键和值向量应用相同的处理，忽略了键和值在注意力计算中的不同角色以及不同令牌的重要性差异。其次，当前的剪枝方法在注意力头之间静态分配内存，未能考虑到不同头和请求之间的动态注意力模式。此外，这些剪枝方法尚未在需要高级推理的工作负载（如编程和数学）上进行全面评估。最后，先前的工作大多忽视了将量化和剪枝统一起来以实现更高服务效率的潜力。

在本文中，我们提出了LeanKV，一个统一的KV缓存压缩框架，旨在解决现有方法的局限性并提高LLM的服务效率。LeanKV引入了三项算法原则：

1. 异质KV：LeanKV使用异质精度压缩键和值向量，键向量以更高精度存储，值向量以较低精度存储。这一设计基于键向量对注意力计算的影响更大的观察，因为它们影响所有令牌的注意力分数，而值向量仅影响单个令牌对注意力输出的贡献。

2. 每头动态稀疏性：LeanKV动态识别每个头和每个请求的关键令牌，并相应地调整内存分配。这一设计基于注意力分数分布在头和请求之间表现出动态稀疏性的观察。在单个请求中，关键令牌的数量在不同头之间变化，而在同一个头中，关键令牌的数量在不同请求之间也会变化。

3. 统一KV压缩：基于异质KV量化和每头动态稀疏性，LeanKV通过在每个头和每个请求上应用混合精度量化和选择性剪枝来统一KV压缩策略，令牌重要性作为统一指标。关键令牌以高精度存储，不太重要的令牌以低精度量化，最不重要的令牌被剪枝。这种集成方法实现了模型准确性和内存效率之间的平滑权衡。

实现LeanKV的算法创新在系统设计上带来了新的挑战，特别是在内存管理方面。与传统的注意力机制采用跨注意力头的统一内存分配不同，LeanKV需要为每个请求动态分配每头内存，显著增加了系统复杂性。此外，LeanKV必须灵活配置内存以适应以不同精度级别存储或完全剪枝的令牌数量，进一步复杂化了内存管理。

为了解决这些挑战，LeanKV引入了两项新技术。首先，LeanKV采用统一分页，将内存划分为固定大小的页面，包含以各种精度存储的向量。每个页面根据头和请求所需的精度配置动态解析，同时注意力GPU内核与页面布局共同设计以最大化内存带宽利用率。其次，LeanKV引入了GPU上的内存管理机制，通过在GPU上并行化内存管理并最小化元数据开销来加速内存管理。

我们在vLLM上实现了LeanKV，并在三个主要模型系列上进行了评估，使用了代表性的工作负载，包括需要高级推理的任务（如数学和编程）。实验结果表明，LeanKV在成本-准确性权衡方面优于现有的KV缓存管理方法。具体而言，LeanKV在不损失模型准确性的情况下将KV缓存压缩了3.0倍到5.0倍，并在准确性损失小于5%的情况下压缩高达11.0倍，吞吐量提高了1.9倍到2.5倍，最高可达6.9倍。

2. 背景与相关工作

在本节中，我们描述了基于Transformer的大型语言模型，并回顾了旨在提高LLM服务效率的现有KV缓存管理技术。

2.1 大型语言模型

Transformer架构：大型语言模型（LLMs）主要基于Transformer架构，用于捕捉序列数据中的复杂长程依赖关系。Transformer架构的核心是注意力机制，它允许序列中的每个令牌在构建其上下文表示时权衡其他令牌的重要性。数学上，注意力计算定义为：

为了捕捉令牌之间更广泛的交互，Transformer模型采用多头注意力（MHA），其中每个头使用不同的查询、键和值投影独立计算注意力。不同头的输出被连接并线性变换以创建更具表现力的表示。分组查询注意力（GQA）通过允许一组查询头共享相同的键和值投影来提高MHA的效率，称为KV头。

自回归生成与KV缓存：LLMs的自回归生成过程包括两个阶段：提示阶段，模型计算提示中所有令牌的潜在表示并生成第一个新令牌；生成阶段，模型通过关注提示和先前生成的令牌迭代生成后续令牌。为了避免跨生成步骤的冗余计算，引入了KV缓存来存储所有先前令牌的键和值。然而，KV缓存的大小随序列长度和批量大小线性增长，迅速成为推理吞吐量的瓶颈。因此，高效的KV缓存管理对于缓解内存瓶颈和提高LLM服务效率至关重要。

2.2 KV缓存优化

分页注意力：静态KV缓存管理系统为最大可能的序列长度预留内存，导致实际序列较短时内存浪费严重。为了缓解这种低效性，vLLM引入了分页注意力，将KV缓存划分为页面，每个页面包含固定数量的令牌，并根据序列长度的增长按需分配页面。通过以页面粒度管理KV缓存，vLLM减少了内存浪费，并支持更大的批量大小以提高服务效率。

最先进的KV缓存量化方法如Atom和Qserve独立地对每个键和值向量应用此过程。然而，这些统一的量化方法在所有令牌的键和值上应用固定比特宽度，可能不是最优的。它们忽略了令牌的不同重要性以及键和值在注意力计算中的不同角色。

KV缓存剪枝：最近的一些工作探索了KV缓存剪枝以缓解LLM推理中的内存瓶颈，这可以被视为量化的极端情况。这些方法使用注意力分数评估令牌重要性并淘汰不太重要的令牌，从而减少KV缓存大小。然而，这些方法的有效性受到其静态和不灵活的内存资源分配的限制。具体来说，H2O和SnapKV在所有注意力头和层之间均匀分配相同的内存预算。PyramidKV基于经验观察在较低层分配更多内存，在较高层分配较少内存，但仍依赖于静态启发式方法，无法动态调整缓存分配以适应各种工作负载。与KV缓存剪枝互补的另一项工作通过将整个KV缓存卸载到CPU内存并选择性地将重要令牌加载到GPU进行注意力计算来解决GPU内存限制。然而，这些方法并未从根本上减少KV缓存大小，并且由于CPU和GPU之间的数据传输而产生了显著的延迟。

3. 动机

在本节中，我们介绍了支撑LeanKV设计的关键发现及其对KV缓存管理的影响。

3.1 键和值的影响

虽然不同令牌的重要性可以直接从注意力分数中推断出来，但键和值向量在令牌内的作用却不那么明显。然而，检查公式1揭示了它们影响的差异。每个令牌对注意力输出的贡献取决于两个因素：来自softmax的注意力分数和值向量。键向量作为共享softmax分母的一部分，影响所有令牌的注意力分数，而值向量仅影响其各自令牌对输出的贡献。

除了键向量的更广泛影响外，建立由键向量决定的注意力分数相对于值向量的相对重要性至关重要。如果没有这一点，键向量的重要性可能会被值向量的重要性所掩盖。为此，我们将注意力机制重新表述为单位向量的加权和，如公式3所示。这种表述将每个输入令牌的贡献分解为两个部分：单位向量​​，表示令牌在特征空间中的方向；系数，定义为注意力分数与值向量范数的乘积，决定了令牌方向的相对重要性。这使我们能够通过检查这些系数的贡献来评估注意力分数和值向量的相对重要性。

图1展示了Llama3-8B模型在Wikitext数据集前1000篇文章上评估的三个代表性层的平均注意力分数和值向量范数的分布。值得注意的是，注意力分数跨越了七个数量级，远远超过了值向量范数的范围，后者最多只跨越两个数量级。在更大的Llama3-70B模型中也观察到了类似的模式。这种显著差异突显了注意力分数在确定每个令牌对注意力输出贡献中的关键作用。因此，我们得出结论，键向量比值向量具有更广泛和更深远的影响，这促使我们进一步探索以不同精度级别处理键和值向量。

3.2 动态稀疏性模式

现有的KV缓存剪枝方法主要依赖于为LLMs中的每个Transformer层静态分配内存，要么均匀分配内存，要么基于离线分析分配内存。然而，LLMs在注意力头和单个请求之间表现出动态稀疏性模式，静态内存分配无法充分处理。为了评估注意力计算的稀疏性，我们分析了保留大部分信息所需的最少关键令牌数量，具体通过保留目标百分比（例如95%）的总注意力分数。

我们首先研究了跨层的动态稀疏性模式。使用Wikitext数据集的前1000篇文章，我们评估了Llama3-8B模型。图2展示了每层平均关键令牌数量，跨所有KV头聚合，以保留95%的总注意力分数。垂直条表示标准差，捕捉了单个请求之间的变异性。值得注意的是，稀疏程度在不同层之间差异显著，如关键令牌数量的差异所示。

接下来，我们深入研究了单个层内的动态稀疏性。图3展示了三个代表性层中每个KV头的平均关键令牌数量，水平条表示跨请求的标准差。稀疏性模式仍然高度动态：在每层内，关键令牌数量在不同KV头之间差异显著。此外，在单个KV头内，关键令牌数量在不同请求之间也可能大幅变化，如某些KV头的高标准差所示。

这些高度动态的稀疏性模式，跨层、注意力头和单个请求，强调了在每个头和每个请求基础上进行自适应内存管理的必要性。

3.3 统一量化与剪枝

量化通过减少比特宽度来表示特征值来压缩KV缓存，而剪枝可以被视为量化的极端形式：有效地以零比特存储特征值。通过使用注意力分数来衡量令牌的重要性，我们可以通过以逐渐降低的精度级别存储越来越不重要的令牌来统一KV缓存压缩的频谱，最终导致最不重要的令牌被剪枝。

为了研究这种统一方法的可行性，我们进行了一个简单的实验：将重要令牌的键和值向量量化为8位，同时剪枝不重要的令牌。图4比较了FP16基线、8位量化、剪枝以及8位量化与剪枝结合在CNN/Daily Mail摘要数据集上使用Llama3-8B模型的性能。对于剪枝和结合方法，大约20%的不重要令牌被剪枝。ROUGE-1分数（衡量LLM生成响应与参考答案的保真度）在四种方法之间几乎无法区分。然而，8位量化与剪枝结合的方法实现了最低的内存使用量。这种结合方法的效率提升促使我们进一步探索通过引入不同精度级别以适应不同重要性的令牌来统一KV缓存压缩的频谱。

3.4 关键发现与启示

在本节中，我们概述了激励LeanKV设计的关键发现。具体来说，我们证明了在单个令牌内，键向量比值向量具有更大的影响，这促使我们提出以更高精度存储键向量而非值向量。此外，不同令牌的重要性（通过softmax分数体现）在头和请求之间表现出动态稀疏性模式。这些模式要求在每个头和每个请求的基础上进行灵活的内存管理。最后，我们展示了将量化和剪枝统一起来的效果优于单独使用其中任何一种方法。

这些发现激励我们进一步探索KV缓存压缩的频谱，以在每个头和每个请求的基础上实现最大的内存效率。具体来说，每个注意力头应该被允许根据单个请求的稀疏性模式，以不同的精度级别存储令牌，并剪枝最不重要的令牌。此外，在每个令牌内，键和值向量应以不同的精度存储。然而，这种统一方法由于导致不规则的内存布局，给内存管理系统带来了重大挑战，其中每个头和请求的KV缓存由以不同精度存储的不同数量的令牌组成。

4. 统一KV缓存压缩算法

在本节中，我们详细介绍了LeanKV的统一KV缓存压缩算法。该算法基于以下观察：键向量比值向量对注意力计算的影响更大，且注意力分数分布在头和请求之间表现出动态稀疏性。因此，LeanKV采用混合精度量化和选择性剪枝来压缩KV缓存，同时保持模型准确性。

4.1 异质KV量化

LeanKV使用异质精度压缩键和值向量，键向量以更高精度存储，值向量以较低精度存储。具体来说，LeanKV支持以下量化配置：

• K8V4：键向量以8位精度存储，值向量以4位精度存储。

• K4V2：键向量以4位精度存储，值向量以2位精度存储。

这种设计基于键向量对注意力计算的影响更大的观察，因为它们影响所有令牌的注意力分数，而值向量仅影响单个令牌对注意力输出的贡献。

4.2 每头动态稀疏性

LeanKV动态识别每个头和每个请求的关键令牌，并相应地调整内存分配。具体来说，LeanKV使用注意力分数来衡量令牌的重要性，并根据重要性动态分配内存。对于每个头，LeanKV保留一个最近窗口，其中所有令牌以高精度量化，其余令牌根据重要性排序。最不重要的令牌被剪枝，次重要的令牌以低精度量化，最重要的令牌以高精度量化。

4.3 统一KV压缩

基于异质KV量化和每头动态稀疏性，LeanKV通过在每个头和每个请求上应用混合精度量化和选择性剪枝来统一KV压缩策略。具体来说，LeanKV将KV缓存分为两部分：高精度部分​和低精度部分​。关键令牌以高精度存储，不太重要的令牌以低精度量化，最不重要的令牌被剪枝。

4.4 提示阶段与生成阶段

5. 系统设计

在本节中，我们介绍了LeanKV系统的设计，该系统利用提出的统一KV缓存压缩算法来提高LLM推理效率。我们首先讨论了压缩算法对系统性能带来的挑战，然后提出了解决这些挑战的新技术。

5.1 挑战

灵活分页的需求：在分页注意力中，所有令牌以相同精度存储，使得固定页面格式足够。然而，异质KV和混合精度量化引入了不同令牌之间的不同精度级别，使得固定页面格式不足。具体来说，固定页面格式需要保守地为所有令牌分配高精度槽，而不管它们的实际精度如何，导致大量内存浪费。假设内存分配为最高精度K8V4，一个具有K4V2精度的令牌将浪费50%的内存。此外，这种固定页面格式会导致内存访问不对齐，从而阻碍内存带宽利用率并降低整体计算效率。为了充分发挥异质KV和混合精度量化的优势，我们需要一种更灵活的分页设计，能够动态适应不同的精度级别，同时确保紧凑的内存使用和高效的内存访问。

增加的内存管理开销：LeanKV中的每头KV缓存压缩算法在内存管理中引入了显著的延迟和元数据开销。在分页注意力中，内存管理的复杂性为O(#requests)，而不管头的数量如何，因为它跨头均匀分区内存并为每个头分配相同的页面ID。相比之下，LeanKV必须处理不同头之间不同数量的高精度和低精度令牌，这将复杂性增加到O(#requests×#heads)。如果不仔细优化，这种开销可能会掩盖缓存压缩带来的性能提升。此外，为了跟踪高精度和低精度令牌，一种直接的方法是使用单独的数据结构来存储每种精度，从而导致元数据开销增加。

5.2 GPU上的内存管理

对内存管理过程的详细分析揭示了并行化和加速的机会。具体来说，统一压缩算法的内存管理包括两个阶段：规划和协调。在规划阶段，每个头独立确定其内存分配需求。随后的协调阶段同步这些每头需求并将它们映射到GPU的物理内存。例如，假设两个注意力头分别需要16和32字节，如规划阶段所确定，且所有GPU物理内存可用。然后在协调阶段，物理内存字节0到15分配给第一个头，字节16到48分配给第二个头。内存管理复杂性增加的主要来源在于规划阶段，该阶段是完全并行的，非常适合GPU的并行计算能力。协调阶段的并行化更具挑战性，因为它需要跨头同步。然而，我们观察到该阶段可以通过并行前缀和有效地并行化。

因此，我们设计了一个GPU上的内存管理系统，以高效处理每头动态内存分配和回收。该系统引入了三个GPU驻留数据结构，共同解决上述挑战：（1）统一分页以实现灵活分页，（2）循环空闲页面列表以高效执行内存管理中的协调阶段，（3）双向页面表以最小化跟踪高精度和低精度令牌的元数据开销。

统一分页：我们将可用的GPU内存划分为大小均匀的页面，每个页面在分配时配置为以给定精度存储令牌。每个统一页面由六个段组成：量化键、键的量化元数据、量化值、值的量化元数据、令牌分数和位置。量化元数据包括键和值向量的比例和零点。每页存储的令牌数量根据量化配置进行调整，确保紧凑的内存使用。统一分页允许灵活处理具有不同精度的令牌，同时保持一致的内存布局，有效支持我们的统一压缩算法。此外，通过将键、值及其元数据整合到单个结构中，统一分页增强了数据局部性，并消除了分散查找的需要，从而提高了注意力计算期间的内存访问效率。

循环空闲页面列表：为了促进高效的页面分配和回收，我们使用一个集中的GPU驻留数据结构循环空闲页面列表来管理所有可用的页面ID。该列表使用一对开始和结束指针跟踪空闲页面，指针在到达列表末尾时回绕到开头，保持循环结构。当分配页面时，开始指针递增到下一个可用页面ID；当页面被释放时，结束指针递增以将释放的页面ID添加回列表。

值得注意的是，列表中的可用和不可用区域保持连续，使得内存管理中的协调阶段可以通过并行前缀和并行化。具体来说，在每个头确定要分配或释放的页面数量后，并行前缀和操作为每个头计算相对于开始或结束指针的唯一偏移量，确保列表中的操作不冲突。对于内存分配，每个头从其指定区域中并发检索新的页面ID，开始指针按所有头所需的累积页面数量递增。类似地，对于内存回收，每个头并发地将释放的页面ID写入其指定区域，结束指针按释放的页面总数递增。

双向页面表：与分页注意力类似，LeanKV为每个注意力头维护一个页面表，将请求映射到物理内存。页面表中的每个条目是为特定请求分配给头的页面ID列表。为了避免维护高精度和低精度页面的单独页面表所带来的元数据开销翻倍，LeanKV引入了一个统一的GPU驻留数据结构双向页面表。该结构有效地支持KV缓存压缩的三种精度级别，如第4节所提出的，最小化元数据开销。在双向页面表的每个条目中，高精度页面ID从列表的左侧增长，而低精度页面ID从右侧增长，动态适应工作负载的精度需求。每个页面表条目的长度由最大序列长度除以每高精度页面的令牌数确定，确保不会溢出，因为低精度页面始终包含比高精度页面更多的令牌。这种统一方法不仅最小化了元数据开销，还消除了基于精度级别的单独查找的需要，从而提高了注意力计算期间的内存访问效率。双向页面表的内存开销最小：例如，在Llama-2-7B上，批大小为128，具有32层和每层32个注意力头，所有双向页面表的总大小仅为64 MB。相比之下，单个请求的KV缓存占用2 GB。

5.3 内存管理工作流程

在本节中，我们深入探讨了LeanKV的内存管理工作流程，展示了所提出的数据结构如何交互以有效支持统一KV缓存压缩。值得注意的是，此内存管理工作流程在每次推理步骤中对批处理中的所有请求和LLM中的所有层执行一次。这种设计选择确保了足够的并行性以实现高效的GPU执行，并分摊了相关的GPU内核启动开销。

提示阶段：图5展示了提示阶段的内存管理工作流程，使用一个包含八个令牌的请求示例。在此设置中，一个高精度页面可以存储两个令牌，一个低精度页面可以存储四个令牌。由于在运行压缩算法之前，每个头所需的高精度和低精度页面的确切数量未知，我们保守地为每个头分配四个统一页面，假设所有令牌将以高精度存储。然后，每个头独立应用KV缓存压缩算法以确定哪些令牌以高精度、低精度存储或剪枝。此步骤构成了规划阶段，在此期间，每个头根据其独特的稀疏性模式计算其特定的内存需求。在此示例中，第一个头使用一个高精度页面，而第二个头使用两个；两个头还需要一个低精度页面。高精度页面从双向页面表的左侧分配，而低精度页面从右侧分配。在规划阶段之后，任何未使用的页面被标记为回收，触发协调阶段。在这里，我们采用基于并行前缀和的方法，如前一节所述，以促进跨头的高效且无冲突的页面回收。

生成阶段：在每个生成步骤中，仅当某个头的高精度或低精度页面已满时，该头才会分配一个新页面，每个步骤最多需要一个额外的页面。每个头独立检查其页面可用性，并在需要时使用基于前缀和的方法并行分配新页面。与提示阶段不同，生成阶段不执行页面回收，因为存储的令牌总数要么保持不变（如果旧令牌被淘汰），要么增加一个（如果没有淘汰）。一旦请求完成，为该请求分配的所有页面都将被回收，释放内存以供传入请求使用。

6. 实现

我们在vLLM上实现了LeanKV，包含4.5K行CUDA/C++代码和9K行Python代码。我们首先概述了LeanKV的架构，然后详细介绍了我们定制的GPU注意力内核，旨在高效支持统一缓存压缩算法。

6.1 LeanKV架构

LeanKV的架构如图6所示。在每个推理步骤中，调度程序尽可能多地将请求批处理到可用的GPU内存中，以最大化吞吐量，并将选定的请求发送给所有工作器。每个GPU托管一个工作器，负责执行模型的一部分。LeanKV利用张量并行性，其中每个工作器处理每个矩阵乘法的一个分区和一部分注意力头。每个工作器包括一个专用的内存管理器，用于管理其分配的注意力头的KV缓存，使用第5节中描述的方法。由于动态稀疏性模式，每个请求的每个头的内存需求各不相同，导致不规则的内存布局。此外，每个工作器包含一个用于模型计算的执行引擎。为了支持统一KV缓存压缩，执行引擎集成了一个KV压缩器和一个定制的GPU注意力内核。在计算键和值向量后，调用KV压缩器使用第4节中的算法压缩它们，并将结果存储在KV缓存中。定制的GPU注意力内核然后使用压缩的KV缓存高效计算注意力输出。

6.2 高效注意力内核

我们开发了一个定制的GPU注意力内核，高效支持每头动态稀疏性和异质KV量化，充分利用内存访问量的减少并将其转化为加速性能。

总体而言，我们为每个CUDA线程块分配处理每个序列的单个注意力头。然而，混合精度量化引入了潜在的负载不平衡，因为高精度令牌比低精度令牌需要更多的带宽，这是注意力的瓶颈。由于统一分页机制将相同精度的令牌存储在一个页面中，我们通过让线程束首先迭代高精度页面，然后迭代低精度页面来缓解此问题，确保线程束内的线程保持负载平衡。每个页面分两个阶段处理：查询和键之间的点积以得出注意力分数，以及值的加权和。这两个阶段遵循不同的计算模式：分别沿特征维度和令牌维度进行归约。为了确保两个阶段的合并和向量化内存访问，我们设计了定制的数据布局和并行化策略，详细说明如下。

7. 评估

我们首先评估了LeanKV的统一KV缓存压缩算法的有效性，并分析了不同模型对KV缓存压缩的敏感性。然后，我们展示了LeanKV在吞吐量方面的改进。

7.1 实验设置

我们在六个模型上评估了LeanKV，涵盖三个系列：Llama2-7B和70B，Llama3-8B和70B，Mistral-7B和Mixtral-8x7B。Mixtral-8x7B采用混合专家（MOE）架构，而其余模型为密集模型。LeanKV的性能在五个基准测试上进行了评估，分为两类。第一类包括需要高级推理和逻辑的任务：GSM8K（数学）、HumanEval（编程）和MMLU（学术水平问答）。第二类包括传统文本处理任务：SQuAD（阅读理解）和CNN-DailyMail（摘要）。吞吐量和延迟在Nvidia L40 GPU上测量，每个GPU具有48GB内存。

7.2 评估统一缓存压缩

我们首先展示了异质KV量化和动态稀疏性的有效性，然后评估了所提出的统一KV缓存压缩的端到端优势。为了减少浮点计算中的噪声和随机性，每个实验重复五次，数据集随机打乱，报告的结果是这些运行的平均值。

评估异质KV：图7展示了异质KV量化（即K8V4和K4V2）与FP16基线在GSM8K和SQuAD上的性能，这两个基准分别代表了两类任务。为了验证我们的直觉，即键向量比值向量具有更显著的影响（第3.1节），我们还评估了异质KV的镜像配置，即K4V8和K2V4，其中值以比键更高的精度存储。

K8V4在两个基准测试中保持了FP16基线的准确性，而其镜像配置K4V8导致了明显的性能下降，特别是在需要推理的GSM8K基准测试中。K4V2和K2V4之间的性能差距更为显著。在GSM8K上，K4V2保留了至少77.8%的FP16性能，而K2V4将准确性降低到几乎为零。在SQuAD上，K4V2实现了94.7%到99.4%的FP16性能，而K2V4仅实现了28.2%到60.0%。这些结果证实了键向量比值向量起着更关键的作用，并证明了异质KV在保持FP16基线准确性方面的有效性。

评估动态稀疏性：我们评估了每头动态稀疏性（第3.2节）在识别关键令牌方面的有效性，将其与H2O中使用的静态稀疏性方法进行比较，后者为所有注意力头分配相同的内存预算。结果如图8所示，我们使用Mistral-7B模型作为示例，所有令牌以FP16存储。

对于需要推理的工作负载（如GSM8K和HumanEval），动态稀疏性显著优于静态稀疏性，随着剪枝令牌比例的增加，性能差距扩大。然而，对于传统文本处理任务，动态稀疏性的优势不那么明显：在SQuAD上略优于静态稀疏性，在CNN-DailyMail上表现相似。

SQuAD和CNN-DailyMail之间的差异可以归因于它们对令牌剪枝的敏感性不同。SQuAD（阅读理解）需要从提示中精确检索特定单词，使其对剪枝关键令牌更敏感。相比之下，CNN-DailyMail（摘要）本质上涉及丢弃不太相关的信息以生成简洁的摘要，因此对令牌剪枝的敏感性较低。总之，动态稀疏性在需要推理的任务和涉及精确信息检索的文本处理任务上优于静态稀疏性。

评估统一KV缓存压缩：最后，我们评估了第4节中描述的统一KV缓存压缩的端到端性能。对于每个基准测试，我们在高精度阈值ThTh​（离散值为{0.5%,1%,5%,10%,20%}）和低精度阈值Tl​（离散值为{0%,0.1%,0.5%,1%}）上进行网格搜索，选择在不同性能约束下最小化KV缓存内存预算的配置。考虑到浮点计算中固有的随机性，如果配置的平均性能在FP16基线的三个标准差范围内，则认为其性能中性。

表1报告了LeanKV的KV缓存内存使用情况，相对于vLLM进行了归一化，涵盖了所有评估的模型和基准测试。结果展示了两种配置：性能中性，确保输出质量不下降；性能下降小于5%，这是先前工作中常见的目标，能够在更激进的内存节省下保持可接受的质量。在所有模型和工作负载中，LeanKV在性能中性的情况下，将KV缓存内存预算压缩至18.9%到28.9%，在性能下降小于5%的情况下，内存预算压缩至7.6%到20.6%。

此外，图9展示了LeanKV在推理密集型基准测试GSM8K和文本处理基准测试SQuAD上的Pareto性能-内存曲线。结果与FP16基线、传统的8位量化（K8V8）和4位量化（K4V4）以及仅剪枝技术H2O进行了比较。对于GSM8K，K8V8保持了FP16基线的性能，而K4V4导致了2.8%到7.3%的明显性能下降。相比之下，LeanKV在仅使用19.8%到23.1%的内存预算的情况下，实现了与FP16相当的性能，甚至比K4V4所需的25%内存更少。此外，LeanKV提供了更灵活的内存-性能权衡，在性能下降小于5%的情况下，仅使用9.4%到16.3%的内存。相比之下，在类似的内存约束下，H2O的准确性几乎为零。在SQuAD基准测试上也观察到了类似的趋势：LeanKV在内存使用少于K4V4的情况下，匹配了FP16和K8V8的性能，而K4V4导致了明显的性能下降，并且在可比的内存预算下优于H2O。

值得注意的是，在所有评估的模型和基准测试中，LeanKV在使用K8V4-K4V2配置时实现了性能中性，其中关键令牌以K8V4精度存储，非关键令牌以K4V2精度存储，且不进行剪枝。对于性能下降小于5%的设置，最佳配置因模型而异。对于大型70B模型，性能下降小于5%是通过更激进的K4V2-K4V1量化实现的，而不进行剪枝，因为即使在K8V4-K4V2的基础上进行剪枝，也会显著降低性能。相比之下，对于较小的模型，性能下降小于5%是通过在K8V4-K4V2的基础上进一步剪枝最不重要的令牌实现的。这些在低内存情况下的不同配置也在图9中可视化，其中每个LeanKV数据点都标注了其相应的配置。这种差异突显了模型对KV缓存压缩技术的不同敏感性，强调了需要一个能够适应多样化场景的统一框架以最大化效率。

总之，LeanKV提供了优于传统方法的性能-内存权衡。它在保持FP16和8位量化性能的同时，使用的内存甚至少于4位量化。此外，LeanKV在低内存情况下（10%到20%的内存）提供了平滑的权衡，而仅剪枝方法在这些情况下会经历显著的性能下降。

7.3 对KV缓存压缩的敏感性

我们研究了模型对KV缓存压缩的不同敏感性，以Llama模型家族为例，因为它们的训练数据和方法相似。图10展示了Llama2-7B、Llama3-8B和Llama3-70B在剪枝令牌比例增加时相对于FP16基线的归一化性能下降。在这些模型中，Llama2-7B表现出最少的性能下降，即使与参数规模相当的Llama3-8B相比也是如此。这种弹性可以归因于Llama2-7B的多头注意力（MHA）架构，与Llama3-8B和Llama3-70B使用的更高效内存的分组查询注意力（GQA）架构相比，它在KV缓存中引入了更大的冗余。Llama3-70B经历了最显著的性能下降，表明较大的模型在令牌之间更均匀地分配注意力，使其对剪枝更敏感。

图11进一步展示了K4V2-K4V1配置在存储为K4V1的令牌比例增加时的归一化性能下降。Llama3-70B模型对低比特量化表现出最大的弹性，即使60%的令牌存储为K4V1，仍能保持接近基线的性能。相比之下，较小的模型即使所有令牌存储为K4V2，也显示出明显的性能下降，突显了它们对激进量化的更大敏感性。

这些观察结果证实了之前提到的两个关键趋势：较大的模型和采用更高效内存网络架构的新模型对剪枝更敏感，而较小的模型对低比特量化的负面影响更大。

7.4 评估延迟和吞吐量

在本节中，我们对LeanKV的性能进行了全面评估，将其与包括vLLM和H2O在内的最先进系统在各种模型和配置下进行了比较。所有实验均在一台配备八块NVIDIA L40 GPU的机器上进行。对于适合单个GPU内存的模型（如Llama2-7B和Llama3-7B），我们使用一块GPU。对于较大的模型（如Llama2-70B和Llama3-70B），我们将计算并行化到四块GPU上。我们的评估重点在于延迟分解、注意力内核加速和端到端吞吐量，以展示LeanKV的效率和可扩展性。

内存管理开销：图12展示了LeanKV在序列长度为1024令牌时一次推理步骤的延迟分解，评估了不同批量大小下的情况。内存管理开销非常低，在提示阶段占总延迟的不到0.2%，在生成阶段占不到0.9%。模型执行在提示阶段和生成阶段均占主导地位，分别占总延迟的96-97%和92-93%。值得注意的是，随着批量大小的增加，模型执行所占的时间比例略有上升，反映了LeanKV的可扩展性。相比之下，在vLLM中，模型执行在提示阶段占总延迟的97%，在生成阶段占95%。这些结果表明，LeanKV有效地最小化了内存管理和KV压缩的延迟贡献，确保这些辅助过程不会成为性能瓶颈。

注意力内核加速：图13展示了LeanKV的定制注意力内核在不同量化配置下相对于vLLM的加速比。LeanKV实现了与KV缓存大小减少成比例的近线性加速。例如，使用K8V8配置，KV缓存大小相对于FP16数据类型减半，理论加速比为2倍，而LeanKV实现了1.7倍。轻微的差距主要是由于访问量化元数据和执行反量化的开销。此外，LeanKV在较长序列上实现了更大的加速，表明其带宽优化技术随着序列长度的增加而变得更加有效。

端到端吞吐量：图14展示了LeanKV与vLLM和H2O的端到端吞吐量和实现的批量大小。我们在从ShareGPT数据集中采样的1000个序列上进行了实验，生成长度设置为4096令牌。LeanKV的压缩阈值基于最具挑战性的数据集MMLU进行选择，使其成为最保守的选择。在无准确性损失的情况下，LeanKV在各种模型上实现了相对于vLLM的2.0倍到2.5倍的吞吐量加速。当允许5%的准确性损失时，LeanKV实现了更高的加速，范围从3.0倍到6.9倍。这些改进与LeanKV支持的更大批量大小直接相关，得益于其KV缓存压缩技术。例如，在Llama2-7B模型上，LeanKV在无准确性损失的情况下支持62的批量大小，在5%准确性损失的情况下支持177的批量大小，远远超过vLLM的22批量大小。增加的批量大小使LeanKV能够最大化GPU计算容量的利用率，从而实现更高的吞吐量。Llama3-7B的加速比相对于Llama2-7B较小，因为它使用了GQA，将其KV缓存大小减少了4倍。因此，vLLM已经支持相对较大的51.5批量大小，缩小了LeanKV可以利用的批量大小优势。尽管如此，LeanKV在Llama3-7B上仍实现了高达3.4倍的吞吐量提升。

值得注意的是，H2O支持的批量大小甚至比vLLM更小，因此吞吐量更低。这是因为H2O在提示阶段消耗了额外的内存用于计算压缩指标的中间结果，限制了可以支持的批量大小。相比之下，我们设计了一个与FlashAttention兼容的内核，使用Triton避免了中间结果带来的额外内存开销，并实现了与vLLM相当的性能。

8. 结论

我们提出了LeanKV，一个统一的KV缓存压缩框架，旨在通过三项关键技术提高LLM的服务效率：异质KV、每头动态稀疏性和统一KV压缩。为了支持这些创新，LeanKV引入了系统优化，包括统一分页和GPU上的并行内存管理。我们的实验表明，LeanKV在不影响模型准确性的情况下，实现了1.9倍到2.5倍的吞吐量提升，并在最小准确性损失的情况下实现了高达6.9倍的吞吐量提升。