KVSharer：基于不相似性实现跨层 KV Cache 共享

一、背景

本文中我们介绍一种最新的 KV Cache 共享论文 KVSharer，与之前常见的层内共享不同，KVSharer 主要关注跨层共享，并且是整个层的共享。

对应的论文：[2410.18517] KVSharer: Efficient Inference via Layer-Wise Dissimilar KV Cache Sharing

对应的代码库：​​​​https://github.com/yangyifei729/KVSharer/tree/main​​

二、摘要

LLM 推理过程中对 GPU 内存的需求不断增加，而其中 Attention 的 KV Cache 占据了超过 80% 的空间。当前，大多数现有的 KV Cache 压缩算法主要集中在单个 Transformer 层内的压缩，而较少有研究涉及层间压缩（PS：其实我们之前介绍过的 Character.AI 方案和 MixAttention 都是层间共享）。

本文中，作者提供了一种即插即用的方案，称为 KVSharer，通过在层间共享 KV Cache 来实现层间压缩。作者发现一个反直觉的现象：共享相似性较低的 KV Cache 更能保持模型性能（PS：这里实现时其实有个约束，要同时避免对 LM head 之前的 Hidden State 的影响过大）。实验表明，KVSharer 能够在减少 30% 的 KV Cache 计算，以及内存开销，而对模型性能影响不大，同时还能实现至少 1.3x 的生成加速。此外，作者实验表明，KVSharer 与现有的层内 KV Cache 压缩方案兼容，两者结合可进一步节约内存。

三、方案

3.1 概述

如下图 Figure 2 所示为本文方案的示例，其主要包含几个步骤：

• 校正集上推理，并记录 KV Cache。
• 计算任意两层的欧式距离。
• 排序，按相似性从低到高排列。
• 从排序的组合中依次筛选可以共享的层。
• 在整个后续推理中保持离线搜索的共享策略。​

3.2 搜索策略

如下图 Algorithm 1 为详细的搜索算法，其大体包含两个阶段：

KV Cache 相似性计算和初始化（1-4）：

• 首先使用校准集所有样本进行推理，获得所有 Layer 的 KV Cache。
• 在每一层，对所有样本的 KV Cache 求平均。
• 将每一层的 KV Cache 拉平为一维 Embedding。
• 计算任意两层 KV Cache Embedding 的欧式距离（距离越大越不相似），构成 S。
• 按照距离降序排列（距离越大，越不相似），得到 R。

共享策略搜索（5-18）：

• 初始空的共享策略 Z 和空的共享 Layer 个数 P。
• 依次遍历排序后的候选共享层 R：

将当前候选 r 添加到共享策略 Z。

替换候选共享策略 Z 中的所有相关层（保留靠近输入的层，靠近输出的层共享靠近输入的层），然后使用校准集验证最后一层 Hidden Stage 和未替换时的相似度。

如果相似度小于阈值 T，表示替换后影响较大，则从 Z 中删除当前后续 r。

如果相似度大于等于阈值 T，则可以作为候选。

如果候选集容量满足 C，则直接结束。​

PS：需要说明的是，以上是基于贪心搜索策略，这种逐步判断并添加的方式往往只能找到局部最优解，并不一定是全局最优的。要想找到全局最优，通常需要进行全局搜索，比如穷举所有可能的共享组合，或者产生更复杂的搜索算法，但是往往会大幅增加计算成本，因此这种贪心搜索是计算效率和效果的折衷。

3.3 推理

如下图 Figure 3 所示，在推理阶段直接使用离线阶段搜索出来的共享策略，可以看出其实是会存在多个 Layer 共享一个 Layer 的。但是图中似乎有问题，论文中并没有讨论共享的传递性，比如Layer 2 和 3 共享 Layer 1，而 Layer 4 共享 Layer 3，Layer 5 共享 Layer 4，那岂不是 Layer 2，3，4，5 都共享 Layer 1？

如下图所示为代码中的注释（KVSharer/llama_real_share/modeling_llama_kvsharer.py#L364-L371）：

四、实验&结果

4.1 精度

如下图 Table 1 所示，作者测试几个常见模型 LLaMA2-7B/13B（PS：如果有 LLaMA3 的结果更有说服力）、InternLM2-7B/20B（中英文能力不错）和 Mistral-7B 在一些主要任务上的精度。其中 Layer 表示实际计算的层数，Layer 越小，表示共享的越多。Percent 表示共享后的平均精度相比原始平均精度的比例。可以看出，共享越多效果越差，并且各个模型表现各不相同，比如在 LLaMA2-13B（40->30） 和 Mistral-7B（32->24） 在共享 25% 时，平均精度就不到 90% 了。

其实，从困惑度（PPL）上来看，影响也是挺大的，尤其是与 H2O 和 PyramidInfer 这些层内压缩结合后：

4.2 推理

如下图 Table 2 所示，作者以 LLaMA2-13B-Chat 模型为例，对比了不同方案对 Memory 的开销以及推理的加速比。其中的 KVSharer（25%）表示 25% 的压缩比，也就意味着 40 层只计算 30 层的 KV Cache。（PS：这里的结论很奇怪，25% 的压缩比理论上最多节约 25% 内存，最多吞吐提升 25%？更何况 Q 和 Attention 还要正常计算，而实际上 1024+4096 时节约了 36% 的内存，吞吐提升 1.53x。）

4.3 消融实验

如下图 Figure 6 所示，作者实验表明：共享相似性较低的 KV Cache 更能保持模型性能。

如下图 Table 3 所示，使用 Wikipedia 或 BookCorpus 作为校准集对精度的影响不是特别大：

如下图 Table 4 所示，随机共享相比本文的 KVSharer 会导致精度下降比较多，证明本文方法的有效性：

如下图 Table 5 所示，KVSharer 在几个模型的 Base 模型和 Chat 模型上的精度影响比较类似，证明了方法的通用性。（PS：不过 PPL 确实影响挺大的）

五、参考链接

• ​​https://arxiv.org/abs/2410.18517​​
• ​​https://github.com/yangyifei729/KVSharer/tree/main​​​​​​​
• https://github.com/yangyifei729/KVSharer/blob/main/llama_real_share/modeling_llama_kvsharer.py#L364-L371

本文转载自 ​​AI闲谈​​，作者： AI闲谈