共享 KVSharer 基于不相似性成功跨层 KV Cache

一、背景

本文中咱们引见一种最新的 KV Cache 共享论文 KVSharer，与之前经常出现的层内共享不同，KVSharer 重要关注跨层共享，并且是整个层的共享。

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

对应的代码库：​​​ ​​ ​

二、摘要

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 所示为本文打算的示例，其重要蕴含几个步骤：

3.2 搜查战略

如下图 Algorithm 1 为具体的搜查算法，其大体蕴含两个阶段：

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

共享战略搜查（5-18）：

将候选 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 确实影响挺大的）

五、参考链接

本文转载自​​，作者：