大语言模型（LLM）中大数据的压缩存储及其重要性

在大型语言模型（LLM）中，KV Cache（键值缓存）的压缩方法及其重要性。

为什么要压缩KV Cache？

计算效率：在生成文本的过程中，每个生成的token都需要与之前所有的token的键值（K和V）进行交互。如果不使用KV Cache，则每次生成新token时都需要重新计算，这将导致巨大的计算量。
显存消耗：KV Cache虽然减少了计算量，但随着token数量的增加，它占用的显存也会增加。显存是有限的，因此需要对KV Cache进行压缩。

KV Cache的分类及方法

生成阶段压缩：大多数方法集中在生成阶段的KV Cache压缩，以支持长文本的生成或加速计算。这些方法通常采用驱逐策略，保留重要的KV，舍弃不重要的。
ClusterKV：通过聚类KV Cache中的token，并使用聚类中心来判断重要性，减少计算量同时保持精度。
LeanKV：根据每个attention head的需要保留不同的KV，分为高精度和低精度保留，动态调整以压缩KV Cache。

输入prompt的KV Cache压缩

SnapKV：针对输入prompt太长导致的问题，通过识别输入中的模式来压缩KV Cache。它假设这些模式在生成过程中是一致的。
ZigZagKV：在SnapKV的基础上进行优化，提出不应该给每个attention head分配相同大小的显存，而应根据每个head所需的最少token数来分配。

ZigZagKV的空间分配

问题：给模型的每一层分配相同大小的显存来存储KV Cache是不合理的，因为不同层的LMBA（维持attention score所需的最小KV数）差异很大。
解决方案：提出了一种按比例分配显存的方法，并为每个层保留一个最小大小，以避免某些层分配到的空间过小。

效果

文章最后提到，ZigZagKV在各种测试中表现最佳，特别是在大海捞针测试中。

总结

在大型语言模型中，如何通过压缩KV Cache来提高计算效率和减少显存消耗。不同的压缩方法针对不同的场景和需求，而ZigZagKV方法在保持模型性能的同时，有效地解决了显存分配的问题。

source：ZigZagKV，KVCache80%的压缩率可以达到近乎无损的效果

为什么要压缩？

输入Prompt和decoding阶段生成的每个token都会产生对应的K和V。每次生成一个新的token的时候，当前的需要和历史上所有的K和V进行交互得到结果，所以我们一般把已经计算出来的K和V都保存起来供后面的token使用，保存的结果就叫做KV Cache。

如果不使用KV Cache来保存历史结果，那么每一个token生成的时候都需要重新计算prompt和之前生成的tokens的K和V。显然，KV Cache减小了计算量。但是，它增加了显存的消耗。随着token的数量增多，需要越来越多的显存。而我们的显存显然不是无穷大的，这就要求我们需要做KV Cache的压缩。此外，随着KV的增多，每生成一个token的计算量也越来越大(因为需要和所有的K和V都进行计算)。这进一步增加了KV Cache压缩的需求。

KV Cache分类以及方法

KV Cache的内容来源于两个方面：

输入prompt；
生成的token。

之前的大多数文章里面的KV Cache方法都集中在处理生成阶段压缩KV Cache，以此来支持生成长文本或者在生成过程中加速计算。这一类KV Cache压缩的方法大多采取驱逐策略，保留系统认为重要的KV，舍弃那些不重要的。其中重要性的判断非常重要，大部分的系统都是对于每一层按照attention score的大小，保留top-B的结果。

前面的文章：

ClusterKV，接近无损的KV Cache压缩方法
LeanKV, 接近无损的5倍的KV Cache压缩以及2.5倍的吞吐提升

介绍了两种效果都不错的KVCache压缩方法。其中ClusterKV的重要观察点是对于一个q重要的所有token对应的K之间的cos距离都比较小。另外就是KV Cache里面某个token的重要性是动态变化的。基于这两个观察，ClusterKV把KV Cache里面的token进行聚类，使用聚类的中心来做重要性判断，如果一个类被判断为重要那么就把这一个类里面所有的token都拿来做计算，这样既减少了计算量还保持了精度。而KVCache并不常驻显存，这样就可以减少显存使用。这篇文章的方法其实是显著不同于大多数采用驱逐策略的方法，因为每个KV都没有完全被驱逐，都是可重入的。

LeanKV也是一篇做得非常细致的文章。核心观察是每一层的每个attention head需要保留的重要KV不一样，并且具有稀疏性。基于这个观察，它不以layer为单位，而是以attention head为单位，去保留对应的重要KV。保留的KV又会分为高精度保留、低精度保留。每次新生成了KV的时候，每个attention head会去计算是否保留新的K和V，如果保留，是否需要驱逐之前的。用这种动态的方法，达到压缩KV又保持精度的效果。

另外一类文章，以SnapKV为代表，研究的是另外一个问题：输入太长，压缩输入prompt的KV Cache。这里的压缩并不是要让生成的token很长，而是不压缩的话可能输入就已经把显存撑爆。或者即便显存还可以，每一步计算也会因为需要交互的token太多而非常慢。所以压缩输入的KV解决的是这两个问题。

ZigZagKV其实是在SnapKV上继续做优化，也就是说，它也是输入prompt的KV Cache压缩算法。它的核心改进点在于不应该给每个attention head分配同样多的显存，而应该根据每个head可以保留绝大部分信息需要的token数来进行分配(SnapKV给所有的layer的head是平均分配的)。

ZigZagKV的一些新的观察

我们首先定义MBA(MinimumBudget size required to maintain 90% of the totalAttention score)，其实就是维持90%的attention score需要的最小KV数：

其中，是的第i个元素。是当前KV Cache里面的token数量。

由于每个attention head都可以计算出一个MBA值，我们把一层里面所有的attention head的MBA的均值叫做LMBA(LayerMinimumBudget size to maintainAttention score)：

LMBA越高也就意味着需要更多的token来维持精度，也就意味着需要分配更多的显存来存储KV Cache。

我们分别计算了Mistral和LLaMA在WikiMQA数据集上的LMBA，结果如下：

可以看出，不同的模型的对应层的LMBA差距比较大；同一个模型不同层的LMBA差距也很大。

刚才我们分析是基于保持attention score，我们再以输出的相似性这个维度分析一下需要保留多少的KV才能让输出类似。我们定义LMBO(Layer-wiseMinimumBudget forOutput):

其中是用所有的KV得到的输出，是用一部分得到的结果。

我们再次发现，同一个模型不同层在保持输出基本一致的限制下需要的KV数量差距很大。

基于上面的两个观察，我们得出的结论就是给模型的每一层分配同样大小显存来存储KV Cache是不合理的。

ZigZagKV空间分配

为每一层分配多大的空间存储KV Cache是合理的呢？

一个直观的想法就是按照比例来分配，大层多份，小的少分。

上面公式里面的就是每一层按照LMBA计算得到的比例，B是单层的Budget，是模型的L层总共的budget，是第层分到的大小。

这么分其实有一个问题，就是如果某一层的LMBA特别大的时候可能导致某些层分配到的空间特别小，也就意味着这些层基本无法保存任何KV。一个层基本没有KV Cache，效果自然保证不了。所以，我们给每个层保留一个最小大小，分配公式改为

这就是ZigZagKV的主要改进点。ZigZagKV和SnapKV都假设给prompt的KV Cache空间是，SanpKV给每一层均匀分配B空间。而ZigZagKV自适应地分配空间。

解决了每一层KV Cache的空间分配问题，还有一个问题就是要保存那些K和V。这里需要先介绍SnapKV。

SnapKV

如果我们输入的prompt有16K，那么如果存储对应的所有KV无论对于显存还是计算都有极大的压力。那显然，我们需要压缩prompt生成的KV。为了压缩输入KV，SnapKV有一个非常重要的观察：

Pattern can be identified before generation and is consistent during generation

什么意思呢？其实就是说，输入里面哪些重要可以使用输入prompt本身就可以决定。

作者把输入128个token为单位划分window，取靠近尾部的20个window来做计算。prompt前面的部分叫做prefix, 20个窗口叫做observation window。对于prefix里面的任意一个token，假设叫做。我们取20个window里面的某一个window，假设叫。的每一个元素和都会有交互，因为在的元素前面。那么，的每个元素和都会计算出一个attention score。这样，对于窗口，会有128个attention score。取这128个attention score的均值作为的重要性值。

这样，对于窗口，prefix的每个token都可以计算出重要性值。注意，我们计算的时候是按照attention head为单位计算的。计算完之后就可以得到。其中是attention head的数量；是一个窗口内token的数量，这里其实就是128；是prompt除了观察的20个窗口的token之外其他的token的总数。

对于每个attention head，prefix里面每个token和窗口内的每个token都有一个attention score，把它们取平均得到最后的score。这是一个的矩阵。然后每个head可以根据平均score选出重要的topK个token。

也可以根据设定的K值选出topK重要的token，其中，代表压缩率。

在生成一个token的时候，对于当前step的q，它和prefix里面的所有token都会有交互。对于模型的每一层，可以计算出的attention score。我们设定一个阈值，我们认为大于这个阈值的attention score对应的token对于当前生成的token来说都是重要的。以此来选择出生成阶段prefix中的重要token。

我们比较窗口选择出来的重要token和生成阶段选择出来的重要token的重叠率，也就是论文所说的Hit Rate。

用数学公式表达就是：

其中。公式稍微有点复杂，更多还是看我的描述吧，文字容易懂点。

通过比较20个窗口选择的重要token和生成阶段实际选择出来的重要token，我们发现最后一个窗口选择出来的和生成阶段选择的是高度重合的。也就意味着，Pattern can be identified before generation。

在生成了若干个token之后，我们把生成的token也像输入一样分成4个窗口，每个窗口包含128个token。然后使用这四个窗口去选择prefix中的重要token。最后发现，选择出来的和prompt里面最后一个窗口选择的是高度吻合的。

这也就是开始所说的Pattern is consistent during generation.

SnapKV基于这两个观察，放心地用prompt里面最后一个窗口来压缩输入prompt的KV Cache。实际上，作者也观察到只使用上面直接用窗口选择的内容，会造成生成的信息不全。所以，实际上的操作是在计算完均值矩阵之后，对每一行都使用1D Pooling操作。其思想很简单，比如第5个token很重要，那么它相邻的4和6我们也应该选进来，不然5的上下文不全。而要选到4和6最简单的就是pool一下，可以简单理解为跟高斯模糊一个性质。

最后选择的token就是pool之后的topK，再机上最后一个窗口里面所有的token。