KVSharer：通过层间不相似KV缓存共享实现高效推理

📖阅读时长：19分钟

🕙发布时间：2025-02-12

近日热文：全网最全的神经网络数学原理（代码和公式）直观解释
欢迎关注知乎和公众号的专栏内容
LLM架构专栏
 知乎LLM专栏
 知乎【柏企】
公众号【柏企科技说】【柏企阅文】

KV缓存是大语言模型（LLM）高效推理中常用的技术，它存储注意力机制中先前计算的键（keys）和值（values），以便在后续生成过程中复用，从而提高推理速度。然而，它却占据了超过80%的GPU内存消耗。为解决这一问题，本文提出了一种即插即用的方法——KVSharer，它通过在层间共享KV缓存来实现层间压缩。与直观地基于更高相似性进行共享不同，我们发现了一个违反直觉的现象：共享不相似的KV缓存能更好地保持模型性能。

关键贡献

基于共享不相似的KV缓存不会显著降低模型性能这一观察，引入了KVSharer，这是一种无需额外训练的层间KV缓存共享机制，用于高效推理。
KVSharer与当前的层内KV缓存压缩方法兼容，能够在保持良好模型性能的同时，进一步减少内存占用。

KVSharer

一、概述

KVSharer分为两个部分：

首先，对于给定的大语言模型，它会搜索一种共享策略，这是一个列表，指定哪些层的KV缓存应由其他特定层的KV缓存替换。
然后，在后续所有任务的预填充（prefill）和生成过程中，相关层的KV缓存会根据这个列表直接进行替换，从而实现高效推理。

二、策略搜索

策略搜索过程如下图所示：

首先在校准数据集上进行推理，并计算任意两层KV缓存之间的欧氏距离。
然后，按照欧氏距离从大到小对这些KV缓存对进行排序。
随后，依次尝试替换相应的KV缓存，同时确保在替换过程中模型的输出与原始模型尽可能保持一致。

搜索过程可参考以下算法：

（a）准备工作

对于给定的大语言模型 $M$，我们设置目标共享KV缓存层数 $C$。
指定一个校准数据集 $D$，它通常由几个普通句子组成。
使用共享KV缓存的模型和原始模型在 $D$ 上进行前向计算，以获得输出表示，确保这些表示的余弦相似度超过阈值 $T$。

（b）搜索过程

KV缓存相似度计算与初始化（算法第1 - 4行）：
- 首先，使用原始模型 $M$ 在校准数据集 $D$ 上进行前向传递，在每个句子的前向传递过程中保存每层的KV缓存。
- 然后，对所有样本的每层KV缓存求平均，得到每层的平均KV缓存。
- 最后，将每层KV缓存的键和值展平为一维向量，然后分别对键和值求平均，以表示该层的KV缓存。
- 接着计算任意两层KV缓存表示之间的欧氏距离，得到 $S$。再将 $S$ 按降序排序得到 $R$，因为欧氏距离越大，相似度越低。
- 因此，不相似的层对会被优先考虑。然后设置两个变量 $Z$ 和 $P$，用于记录候选的KV缓存共享策略和当前共享层数。
共享策略搜索（算法第5 - 18行）：
- 基于 $R$ 中的值，我们依次选择一对层 $r$ 添加到 $Z$ 中进行共享。
- 共享时，我们用更靠近输入的层替换更靠近输出的层的KV缓存，因为大语言模型中靠近输入端的层更敏感，修改它们可能会导致性能显著下降。
- 通过在前向传递过程中直接用一层的KV缓存替换另一层的KV缓存来应用候选策略 $Z$。
- 使用共享KV缓存的模型和原始模型在校准数据集上进行推理，以获得最后一层的输出表示。
- 对不同句子的这些表示求平均。如果两个模型的平均输出表示之间的余弦相似度超过阈值 $T$，我们保留当前的层对替换 $r$；否则，丢弃它。
- 这个迭代过程持续进行，直到达到预定义的压缩层数 $C$。在迭代结束时，我们通过启发式搜索获得最优的KV缓存共享策略 $Z$。

三、基于KV缓存共享的推理

获得KV缓存共享策略 $Z$ 后，我们将其应用于所有后续推理任务，包括预填充和生成过程。
如下图所示，在前向计算过程中，当根据 $Z$ 需要替换某层的KV缓存时，我们直接从先前计算的层复制KV缓存。随后的计算则按照原始模型的流程进行。

实验

一、结果

下表展示了我们实验的主要结果。“Layer”代表实际计算KV缓存的层数，并展示了模型在不同任务方面的平均值，以及所有任务相对于完整KV缓存的平均得分（以百分比表示）。

Llama2 - 7B和InternLM2 - 7B各有32层，而Llama2 - 13B和InternLM2 - 20B分别有40层和48层。为评估性能，我们以12.5%、25%和37.5%的压缩率对这四个模型应用不同数量的压缩层。
与完整的KV缓存相比，KVSharer在绝大多数任务中的性能下降极小。值得注意的是，当压缩率为25%或更低时，性能接近90%，在某些情况下甚至超过95%。
结果表明，KVSharer有效地保持了模型的整体性能和特定任务性能。
KVSharer在压缩率为25%或更低时能够很好地保持模型性能，甚至在Llama2 - 7B上以12.5%的压缩率时提高了模型的平均性能。

二、策略搜索时间

下图展示了KVSharer对不同模型的搜索时间成本。

结果表明，在模型上搜索共享策略大约需要一分钟或更短时间。这是预期的，因为策略搜索只需要模型在由几个到几十个句子组成的校准数据集上进行几次推理，这个过程在GPU上几分钟内即可完成。

三、与层内压缩的兼容性

下图展示了在Wikipedia数据集上不同压缩率下模型的困惑度（perplexity）。“+H2O”和“+Pyr.”分别表示额外使用H2O和PyramidInfer进行层内压缩。

在KVSharer压缩率为12.5%和25%时，这两种方法只会导致困惑度略微增加。PyramidInfer对困惑度的影响比H2O更低，这是预期的，因为PyramidInfer通常能更好地保持模型性能。
对于不同KVSharer压缩率下的InternLM2和Llama2系列，在压缩率不超过25%时，困惑度通常低于15，甚至可达10以下，这使得模型能够保持良好的生成质量。

四、内存成本和推理速度

下表展示了Llama2 - 13BChat的内存使用量（MB）、预填充时间（s）和生成速度（tokens/s）。“SeqLen.”代表“输入长度” + “最大输出长度”。

当句子长度相对较短，如512 + 32个token时，KVSharer的内存节省效果并不显著，因为当前内存使用主要还是由于模型本身。
随着长度增加，内存减少效果开始显现。当长度达到256 + 2048个token时，内存节省可达30%。
结果显示，在预填充阶段没有加速，但在生成阶段有1.2倍的加速。当长度达到512 + 2048时，在生成阶段可提供超过1.6倍的加速。

结论

本文引入了KVSharer，这是一种专为大语言模型高效推理设计的层间KV缓存共享方法。
实验表明，KVSharer在将KV缓存计算减少30%的同时，能保持主流大语言模型90%以上的原始性能。它还能在生成过程中至少提供1.3倍的加速。
KVSharer可以与现有的层内KV缓存压缩方法集成，以实现更大的内存节省和更快的推理速度。

论文链接：https://arxiv.org/abs/2410.18517
代码链接：https://github.com/yangyifei729/KVSharer

参考文献：
KVSharer: Efficient Inference via Layer-Wise Dissimilar KV Cache Sharing by Yang et al. arXiv:2410.18517

本文由mdnice多平台发布

KVSharer：通过层间不相似KV缓存共享实现高效推理