大模型推理--KV Cache压缩

这篇博客本应完成于24年9月，但是由于入职新公司之后上下班距离较远下班之后无力再动脑，遂拖到现在才完成，而且再看文章内容感觉不少部分都已过时，但还是可以当做一篇综述性文章来读，了解KV Cache压缩的历史。

在之前的博客《大模型推理–PagedAttention》中我们简单提到了KV Cache在超长上下文下过大的问题。为了使得大模型推理能支持超长上下文，我们必须要对KV Cache进行压缩，本博客就介绍一下目前常用的压缩手段。

我们还是引用PagedAttention中给出的例子，看一下KV Cache到底有多大。规模为13B的OPT大模型，单个token的KV Cache就可以达到800KB，计算公式为：2(key和value两个向量) x 5120(embedding维度) x 40(模型层数) x 2(fp16)，而且这还没有考虑batch维度。假设大模型支持的上下文为128K，则prompt部分的KV Cache就会有102G的大小！所以压缩KV势在必行，减少KV Cache就可以让我们在更少的设备上实现对更长上下文的推理，或者在相同上下文长度下让推理的batch更大，从而实现更快的推理速度或者更大的吞吐量。

目前比较常用的压缩方法主要有三种：量化，横向KV Cache抽取和纵向head压缩，如图1所示。量化是其中比较成熟的方法，目前可以普遍支持到int4，但是量化也只是相比原来的fp16压缩了4倍而已。当前的一个热点是如何从长度为N的KV Cache中抽取出对结果影响最大的一些KV Cache出来，避免保留完整的KV Cache。还有一些方法是从优化Multi-head Attention的head数入手减少KV Cache的大小，比较热门的几种方法是MQA、GQA还有最新的MLA。但是这些方法都是在transformer架构下对Attention打补丁的解决方案，可以应付当下，但是未来可能都会被弃用，新的大模型架构Mamba给我们带来了另外一条路，从架构层面就把KV Cache优化掉了，不过Mamba能不能真正代替transformer还得交给时间来验证。下面我们一一展开介绍这些KV Cache压缩方案。

图1 $N\ast h\ast d_k$的KV Cache压缩，每个位置都会有量化，横向选择不同的行叫抽取，纵向只选择一个head叫MQA，选择多个head叫做GQA。

1. 量化

量化是最经典也是最容易想到的KV Cache压缩方案。从大模型推理的角度，不止KV需要量化，权重和激活(每层输出)也都需要量化，一般用WxAxKVx来表示每一部分量化到多少位，目前比较主流的方案可以实现W4A8KV4(QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving)，终极目标肯定是W1A1KV1，但是我觉得这个目标在达到之前transformer应该已经被替换掉了，不过W2A4KV2短时间内还是可以期待一下的。本博客主要关注KV的量化，后续有时间再写一下权重和激活的量化，不过你读过一些论文之后就会发现，它们的量化方案有相似之处，懂了一个之后另一个也能触类旁通。

关于KV Cache的量化，网上有一篇非常好的综述性博客《量化那些事之KVCache的量化》，里面介绍了最新的KV Cache量化论文以及每篇论文的亮点，非常值得读。本博客挑选其中一些比较知名的论文进行详细地展开。

1.1 KIVI

KIVI这篇论文提出了一种比较简单可行的KV Cache压缩方案，可以实现2bit的压缩。论文的最大贡献点在于让大家形成一种共识：Key cache要采用per-channel量化方案而value cache则要采用per-token量化方案。如此做的目的也是发现大模型推理过程中key经常会在某个channel出现异常值，而在value中则比较少出现异常值，如图2所示。从图中很容易看出，key cache的异常值是非异常值的5倍以上而且都集中在特定的channel中。这种情况下采用per-channel量化就可以将误差局限于固定的channel中，而如果要采用per-token量化就会导致精度被异常值吃掉。虽然value cache没有明显的异常值，但是作者发现如果value采用per-channel量化也会导致精度大幅下降，原因在于value要和Attention矩阵相乘，而attention score具有稀疏性，最终的表现就是几个salient token与value相乘。这种情况下如果采用per-channel量化就会导致所有的token都会受到量化误差的影响，而per-token量化则可以将量化误差局限于单个token中。

图2 Llama2-13B和Falcon-7B不同层的key和value cache数值示意图

我们一直在提per-channel和per-token量化，到底何为per-channel何为per-token呢？这个概念也比较好理解，请看图3。给定$l_{prompt}$个token的embedding矩阵，将其与$W_K和W_V$相乘会得到$X_K和X_V。X_K和X_V$按行去量化就叫per-token，因为每一行都是由同一个token的embedding信息生成；按列去量化就叫per-channel，因为每一列的信息包含了所有token的embedding信息，可以被认为是token的某个特定通道。per-token天然满足流式性要求，可以比较好得在大模型推理decoding阶段实现逐token量化；per-channel则相反，它无法实现线性的量化，因为每新来一个token，每个channel的量化都要重新做，所以它是平方级的复杂度。

图3 per-token和per-channel量化示意图

KIVI的巧妙之处就在于它提出了一种取巧的方法解决了key cache的per-channel量化问题：对key cache按行数分组，当达到一定行时就进行key的量化，不足一定行时(论文中称为residual部分)就维持浮点精度，后续的attention矩阵计算就用了定点和浮点的分块矩阵乘实现。由于residual部分很小(论文中为128)，相比整个大模型上下文要小很多，所以对Key的压缩还是非常明显的。另外，在第二部分KV Cache抽取中我们也会提到，Attention矩阵具有一定的局部性，residual部分保持浮点也满足了Attention矩阵的局部性要求，所以KIVI可以做到2bit不明显掉点。

整体上，KIVI的论文实验比较充分，得到的结论也被普遍认可，做法也比较简单，算是一篇不错的论文。但是针对key的异常值处理不是只有per-channel量化一种方案，smoothQuant也是一种可行的方案，将key的异常值迁移到$W_K$中去，或者将两种方案组合起来说不定可以获得更好的结果。论文中也提到，key采用4bit per-token量化掉点也比较小，所以在快速实现时KV都采用4bit的per-token量化也是一种可行的方案。

1.2 QServe

《QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving》是MIT韩松实验室的工作，smoothQuant就是这个实验室提出的，给大模型权重和激活压缩提供了一种比较简单高效的通用方法。QServe算是一篇系统工程论文，不止考虑了KV的压缩，还把权重、激活和KV都考虑进来，实现了一个W4A8KV4的完整大模型压缩方案，在本博客中我们只关注其中的KV压缩方案。
QServe在KV的压缩过程中发现了和KIVI一样的结论：Key的异常值通常分布在固定的通道中，value则没有明显的异常值。不同与KIVI的方案，QServe提出的KV压缩方案叫做SmoothAttention，继承于smoothQuant，将key的异常值迁移到$W_K$中去，然后统一对KV采用4bit的per-token量化。SmoothAttention的基本原理都是基于下述等价变换公式：
$$A=(Q\Lambda )\cdot (K{\Lambda }^{-1}{)}^T$$
由于K中经常出现较大的异常值，我们可以将K进行缩放变小，同时Q等价变大，使得K的压缩难度大大降低。在真实计算中，矩阵不会真正参与运算，而是可以将其迁移到$W_Q和W_K中去。\Lambda$的生成需要采用PTQ后量化方案利用校准测试集进行生成。整体上来说，SmoothAttention和SmoothQuant非常类似。由于SmoothAttention中Q不量化，所以它的迁移难度要小于SmoothQuant，不过它还需要考虑RoPE带来的问题，具体细节可以参看论文。

读了这两篇论文我们就会发现，处理key中异常值的两种方案，将这两种方案结合理论上可以进一步提升KV压缩的准确率，并获得更低的压缩位宽。不过目前的大模型推理论文针对激活和KV中的通道异常值还停留在发现现象的阶段，没有从更深层次去探究原理并从模型训练角度解决该问题。

2. KV Cache抽取

KV Cache抽取的方法在其他博客中也会被称为Attention sparsity或者token dropping(eviction)，但表达的都是从行的角度来对KV Cache进行选择，把重要的行留下来，删掉不重要的行，根据抽取的策略不同又衍生出了非常多的方法，我们介绍其中几个比较知名的策略。

2.1 Window(滑窗)

想到滑窗，那就会衍生出很多不同的滑动策略，这个方向比较知名的工作是streamingLLM，在其论文中对比了三种不同的滑窗策略。streamingLLM选择通过保留开头四个token的KV Cache加最近的L个token的KV Cache可以使得大模型支持的上下文近乎无限长，而且效果也还可以，如图4所示。

图4 原始的Attention与三种不同的滑窗Attention策略

图4中的a表示不对KV Cache进行抽取，但是当输出token的长度超过预训练文本的长度时效果开始急剧变差。b表示只考虑最近的L个token的KV Cache，效果也非常差。c和b类似，但是会将最近L个token进行Attention的重计算，所以速度不太快。d是streamingLLM采取的滑窗策略，它在b的基础上又考虑了最开头的4个token(论文中称之为Attention sinks)，使得效果大幅提升。为啥开头几个token的KV Cache这么重要呢？streamingLLM的研究人员通过实验发现，在Attention计算的过程中，很多层中都有相同的现象：当前的token总是和开头几个token计算出很高的softmax得分，把它们的KV Cache保留下来就可以明显得提升推理效果，如图5所示。

图5 不同层的softmax得分分布

从图5我们可以看出，前两层中每个token与临近的token计算出较大的得分，但是后续的层都是和开头的几个token计算出较大的得分。根据这个发现，streamingLLM才设计出一种开头4个token+临近L个token的KV Cache抽取策略，并取得了不错的效果，图6展示了L为3时的KV Cache抽取示意图。

图6 streamingLLM的KV Cache抽取策略

论文中用来验证Attention sinks结论的测试用例都太短，在真实长上下文下是否有类似的结论其实不太确定(在其他论文中指出，不是所有的大模型都会存在Attention sinks)。另外如果滑动窗口中的L比较大，比如16K说不定也可以达到不错的效果，毕竟大部分时候生成一个主题的内容也不会太长。但是不管怎样，滑窗策略肯定是丢失了很多信息，在大海捞针测试中效果肯定会变差。

2.2 $H_{2}O$

和streamingLLM同一时期还有一个类似的工作，效果也会更好，名叫$H_{2}O$。 在更多LLM对比下，我们就会发现streamingLLM的token抽取策略还是太粗糙了，很难想象单纯靠开头几个token和临近token就可以保持好的效果。$H_{2}O$通过对Attention矩阵更深入的洞察提出了一种更好的token抽取策略，效果相比streamingLLM有比较明显的提升，如图7所示。

图7 streamingLLM和H2O在文档问题回答和摘要生成上的效果对比

$H_{2}O$基于如下几个发现提出了全新的token抽取策略：

1. Attention矩阵的稀疏性。统计不足Attention score最大值1%的比例我们可以获得图8a，这说明Attention矩阵中大部分得分都非常小，去掉这些得分对最终结果也不会产生太大的影响，这个比例接近95%，也即意味着有上限20倍左右的token压缩比；
2. heavy-hitters ($H_2$)的存在。$H_2$表示关键token，去掉它们会导致生成效果的急剧下降。作者发现每个token的累计Attention得分满足幂律(power-law distribution)，也即少数token的得分占据了大部分，如图8b所示，这与上面的结论一致。此外，作者还发现每个单词的共生词也满足幂律(we can see the accumulated attention score of each word (in red dots) have a high correlation with their co-occurrences in the data (gray curve))。由此我们就可以设计一个H2加临近token的抽取策略在大幅降低KV Cache的情况下保证效果；
   

图8 (a)不同LLM每一层Attention矩阵的得分稀疏性 (b)每个token的累计Attention得分满足幂律

3. $H_2$依赖于完整的token列表，但是在推理时这个是无法获取的。为了解决这个难题，作者提出了local $H_2$用来实现快速的token驱逐，具体的过程是在每一次decoding时将之前每一个已生成token的Attention得分累计起来按得分进行驱逐，这样得到的推理效果和全局H2类似，如图8d所示。图9展示了H2的驱逐策略，KV Cache窗口大小为3，通过计算Attention矩阵每一列的得分和进行排序驱逐。
   

图9 H2驱逐策略

$H_{2}O$的作者还对比了$H_2$和临近token在推理时的贡献度，发现两者对保证最终的结果都有价值，$H_2$的作用反而会更大，如下表所示。当只保留临近几个token的KV Cache时效果下降会比较明显，当只保留H2时效果也会有下降，但是会比单纯保留临近token的KV Cache要好，将两者综合起来就可以接近完整KV Cache的结果。

$H_{2}O$最终可以实现在5倍压缩KV Cache的情况下保证推理效果基本不变，但是很遗憾现在几乎没有人用，最根本的原因就在于它依赖于Attention矩阵，而当前最流行的prefilling推理手段FlashAttention偏偏通过分块把Attention矩阵的计算给优化掉了，这就导致$H_{2}O$和FlashAttention不兼容，在FlashAttention成为主流的情况下$H_{2}O$就很难被推广。

虽然$H_{2}O$通过深入研究Attention矩阵取得了较高的KV Cache压缩比，但是不可避免还是会丢失掉一些信息，这应该属于KV Cache抽取类方法的固有缺陷。从经验上来说，用户给定的prompt确实包含很多冗余词，去掉这些冗余对效果不会产生影响。但是每个词都有其存在的意义，少了这些词prompt的表达能力可能会受到影响，比如给定一句话“我 他妈的 快要 饿 死了”，它表达了“我饿”这样一个意思。假定我们的KV Cache缓存窗口只有2，很有可能保留下来的token只有“我 饿”这两个token，虽然保留了该句话要表达的意思，但是缺少了强烈的情感，这可能会使LLM推理产生很大的误差。正如俄罗斯领导人曾经说过的那句话“俄罗斯虽大，但是没有一寸土地是多余的”。

streamingLLM和$H_{2}O$都没有特别提及它们的使用阶段，我理解prefilling和decoding都可以使用这两种方案。为了支持长上下文可以在prefilling阶段使用这两种方法，但是会丢掉不少上下文信息且无法与FlashAttention融合；为了支持无限内容生成，可以用在decoding阶段，但是这样还是会出现针对长上下文产生完整的KV Cache然后再压缩的问题，所以各有利弊。

2.3 TOVA

还有一种方案叫TOVA，我觉得特别适合用在decoding阶段的KV Cache压缩中，思路相比$H_{2}O$还会更简单：在KV Cache到达上限之后，TOVA只是计算当前token的attention得分，丢掉得分最低的KV Cache。虽然算法并没有强调一定要在decoding阶段使用，但是我觉得在decoding阶段比较容易计算attention得分，代码实现上会比较容易。在真正使用时，我们不必关心prefilling阶段生成了多少KV Cache，有多少存多少。但是到了decoding阶段，我们就引入TOVA，实时计算当前token的attention得分抛弃得分最低的KV Cache，如此我们就可以压缩掉decoding阶段生成的KV Cache，保持KV Cache的恒定。

在论文中宣称的结果要好于$H_{2}O$，我觉得这个大概率是成立的，相比$H_{2}O$它算是一种更简洁的压缩方法。论文中采用的方法是把attention得分最低的KV Cache丢掉，但是他们也发现这样一种策略保留下来的还是基本以临近token的KV Cache为主，如图10所示。基于图10作者有三个发现，这些发现和streamingLLM基本能一致：

1. Recency is not all you need：临近token重要，但不是全部；
2. The first token matters：开头几个token的重要性很高，它们的KV Cache会存活很久；
3. Not all tokens are equally kept：有些词的重要性格外高，论文中发现像POS ” $ ) . NNPS \n等token存活的时间也格外久。
   

图10 在每一个decoding阶段保留下来的token id

我个人觉得TOVA不是多先进的算法，但是确实可以考虑在实际中使用，作为一个较为初级的KV Cache压缩策略。

2.4 DMC

顺着streamingLLM、H2O、TOVA的论文读，就会发现NVIDIA还有一篇论文《Dynamic Memory Compression: Retrofitting LLMs for Accelerated Inference》在前者的基础上提出一种更好的KV Cache抽取策略。DMC的策略非常简单，在每遇到一个新的KV之后判断是将它添加到KV Cache中，还是将其加权累加到最后一个KV Cache上，如图11所示。判断的标准就是利用新计算出来的Q和K首位计算sigmoid，所以需要对模型进行重新训练，它是一种亚线性的KV Cache压缩方法。可以看出，DMC不依赖Attention矩阵，所以相比之前的几种方法可以比较好得与FlashAttention兼容，但是重新训模型又导致它可能不会被广泛采用。

图11 DMC算法伪代码

这个方向还存在着大量论文，有的论文通过分析跨层的Attention矩阵特点获得更大的压缩比(pyramidInfer)，有的论文通过聚类的思路来压缩KV Cache(SubGen、SnapKV)，方法五花八门。从信息压缩的角度来讲，KV Cache抽取策略确实会取得成功，因为Attention矩阵有其固有的稀疏性，此外token之间确实有重要和不重要的区别。但是我个人感觉这类方法不好，不是一个非常值得去做的研究方向，不够本质，属于对self-attention的修修补补。鉴于它不可避免会带来一定的效果损失，所以在和主流压缩方案兼容之前都不太可能受到大家的重视。

25年2月份，Deepseek又发布了新的论文，给大家介绍了他们家新出的KV Cache压缩方法NSA(Native Sparse Attention)，我自己粗略看了下应该属于KV Cache抽取这个方向，感觉上是综合了上面好几种压缩方法。通过针对现代硬件的优化设计，NSA（原生可训练稀疏注意力机制）在保证性能的同时显著提升了推理速度，并有效降低了预训练成本，详细细节大家可以看一下新智元的文章《DeepSeek革命性NSA注意力机制问世！梁文锋上阵，长文本推理能力飙升》。

3. Head压缩

3.1 MQA、GQA

MQA全称叫Multi-query Attention(多查询注意力)，是19年发表的一篇论文《Fast Transformer Decoding: One Write-Head is All You Need》中提出的，论文写得比较潦草，也不知道有没有正式发表。论文的作者就是transformer论文八大金刚之一的沙哥(Noam Shazeer)，量子位曾经发过一篇公众号文章《OpenAI公关跳起来捂他嘴：Transformer作者公开承认参与Q*！｜八位作者最新专访》详细介绍了transformer的诞生过程，从文章中我们知道沙哥才是transformer的灵魂人物，文风诙谐幽默，大家感兴趣可以看一下。论文刚出来的时候，没有引起太多关注，直到这两年大模型超长上下文推理开始变成各家主打的一个卖点又开始引起不小的关注。

MQA是对MHA的改进，把原始MHA中KV需要h个head的存储优化成只需要一个head的存储，Q的h个head复用这一个head的KV Cache，如此一个操作就把KV Cache压缩到原来的$\frac{1}{h}$，假设大模型的head数是32，就是原来的$\frac{1}{32}$，压缩效率杠杠滴，而且性能损失非常小。

我们先介绍一下基本的MHA：假设Q、K、V在送到Attention模块计算之前的维度为$N\ast d$， MHA会将维度d的向量分成h组变成一个$h\ast d_k$的矩阵，Q、K、V此时成为了$N\ast h\ast d_k$的三维矩阵（不考虑batch维）。分别将Q、K、V的第一和第二维进行转置得到三个维度为$h\ast N\ast d_k$的三维矩阵。此时的三个矩阵就是具有h个head的Q、K、V，我们就可以按照self attention的定义计算h个head的attention值。可以看出，在MHA中我们需要保留大小为N x d的K和V，MQA就是把保留的KV由$N\ast d变为了N\ast d_k$，也即只保留一个head的KV Cache，h个head的Q共享这一份KV Cache。对应的，计算Q、K、V的三个projection矩阵$W_Q，W_K，W_V$维度也不再相同，$W_Q$还是$d\ast d$，但是$W_K，W_V$则变成了$d\ast d_k$，这也表示MQA压缩了模型的参数量。

我们知道大模型的decoding阶段是一个memory-bound的问题，所以压缩KV Cache之后就可以降低decoding时的通信量，进而可以提升decoding的速度，如下图MQA论文中给出的结果所示。dec阶段的速度相比MHA提升了一个数量级，由46降低到3.8。这个结果还是在head数为8上的结论，当前大模型的head数普遍在32以上，所以加速比会更明显。

图12 MQA相比MHA在推理任务上的速度对比

当我们理解了MQA，再去理解Group-quey Attention(分组查询注意力)就很容易了。GQA于2023年提出，可以算是MHA和MQA的折衷方案，既不想损失性能太多，又想获得 MQA带来的推理加速好处，是MQA的更一般形式，也是目前主流大模型普遍采用的方案。具体思想是，它就是将所有Head分为g个组（g可以整除h），每组共享同一对K、V，也即不是所有Q 的head共享一组 KV，而是采用分组的方式让一定head数的Q 共享一组KV，所以GQA的效果要好于MQA，不过推理速度相比MQA也没有太大差异，GQA的原理图如图13所示。

图13 MHA、GQA和MQA示意图

在GQA论文中有一个速度对比的结果值得我们仔细分析，如下图。模型总共有64个head，所以GQA只有一个group对应的就是MQA，有64个group就对应MHA。可以看出GQA的推理耗时基本上符合指数函数的走势，这就说明持续减小KV Cache的大小就可以持续获得速度提升，直到decoding阶段从访存受限(memory bound)到计算受限(compute bound)为止。这也启发我们要持续对KV Cache进行压缩，不过从效果考虑一般会选择8组，从图14也可以看出8组时的推理速度和1组的推理速度差不了很多(这里其实需要更精确的对比才好，论文中没有给出)。

图14 GQA在不同groups下的推理速度

chatGLM2-6B模型以及PaLM、StarCoder、Gemini都采用了MQA进行推理加速，LLAMA2、LLAMA3都采用了GQA进行推理加速，这两个压缩技巧在比较早期的大模型中经常使用，不过估计从24年5月之后它俩就要被MLA替代了。虽然后续的大模型不太可能会使用MQA，但是MQA还是可以作为一个此类方法压缩KV Cache的下限基准而频繁出现。

3.2 MLA

MLA部分请参考《大模型推理–MLA》。

3.3 MFA

现在的大模型推理技术迭代速度太快了，完全跟不上最新的研究成果。25年一月份阶跃星辰和清华大学又推出了一个MFA(Multi-matrix Factorization Attention)，号称性能比MLA还要好，大家感兴趣可以看一下机器之心上面的文章《阶跃公开了自家新型注意力机制：KV缓存消耗直降93.7%，性能不减反增》。不过MFA目前还没有被大家完全吸收，真实效果有待行业进一步验证，待它被大规模采用了我们再详细介绍。

4. 非transformer大模型

Transformer取得成功的关键就在于其平方复杂度的Attention机制，可以让每一个将要生成的token与之前所有的token建立重要性关系，使得transformer可以捕捉全局注意力。但也正是这个平方级的Attention机制使得KV Cache变成了transformer的负担，当下相当数量的研究工作都在集中在如何压缩KV Cache上。是否存在一种模型可以避免平方级的Attention运算，但是效果还能保持transformer的水平？目前来看全面超越transformer比较难，但是还是存在一些非transformer的模型在特定领域可以达到与transformer抗衡的水平。

4.1 Mamba

Mamba是这几年比较流行的非transformer大模型，FlashAttention作者吹岛的另一杰作。Mamba的整体机制比较复杂，在此给出一个简单描述。

Mamba的底层核心是状态空间模型SSM(State space models)，由于SSM具有线性复杂度，能够根据输入有选择地传播或遗忘信息。Mamba结合selective SSM，提升了推理速度，并在多个领域（如语言和基因组学）取得了优异性能，超越了同规模的Transformer模型。简单描述SSM就是由三个矩阵A、B、C构成的两个公式，如图15所示。

图15 离散型SSM

第一个公式和RNN的计算类似，但是没有用非线性的激活函数给包起来，所以记忆和表达能力都要弱一些，但是可以比较快速得推理。将第二个公式中的$h_k$替换成第一个公式并递归展开可以得到如下公式：
$$y_k=\sum _{i=0}^{k}C{\overline{A}}^{k-i}\overline{B}\ast x_i$$
由于其中三个离散参数$\overline{A}、\overline{B}、C$都是常数，因此我们可以预先计算左侧向量并将其保存为卷积核，这为我们提供了一种使用卷积超高速计算y的简单方法。SSM的好处就是可以像RNN一样快速推理，同时还可以像CNN一样快速训练。

Mamba优化了SSM的几个缺陷：

1. SSM模型无法根据不同的输入数据动态调整矩阵B、C以及步长Δ的值，所以引入了selective SSM将部分关键参数（如B、C、Δ）动态地依赖于输入。这种设计不仅增加了灵活性，还使模型可以对不同输入条件调整状态空间参数。与标准 SSM 不同，此算法允许模型在序列的每个时间步上进行调整，实现时间可变性；
2. 由于B、C这些矩阵现在是动态的了，因此无法使用卷积表示来计算它们(CNN需要固定的内核)，因此，我们只能使用循环表示，如此也就而失去了卷积提供的并行训练能力，mamba采用“并行扫描算法”而非“卷积”来进行模型的循环计算；
3. 将SSM架构的设计与transformer的MLP块合并为一个块，来简化过去的深度序列模型架构，从而得到一个包含selective state space的架构设计，如图16所示。
   

图16 mamba block

Mamba是最近几年比较火的非transformer大模型，在某些任务上确实可以超越同等参数量的transformer模型，但是它在大海捞针测试中要比transformer差很多，说明mamba在记忆上要弱于transformer模型。但是它线性的复杂度使得推理不再依赖于KV Cache，推理速度大大提升，显存占用也大幅减少，持续迭代优化还是有一定潜力替代transformer。

4.2 RWKV

RWKV是另外一种高效的深度学习网络架构，它将Transformer与RNN各自的优点相结合，同时实现高度并行化训练与高效推理，时间复杂度为线性复杂度，在长序列推理场景下具有优于Transformer的性能潜力，而且一直在迭代，25年1月8号已经发布到RWKV-7。
RWKV与其他Transformer变体的复杂度对比如图17所示。可以看出，常规的Transformer复杂度与序列长度T的平方成正比，线性Transformer将平方级的复杂度转移到隐藏维度上去也大大降低了复杂度，AFT-local通过点乘运算可以进一步降低复杂度，RWKV则更近一步，实现了完全的线性复杂度。

图17 RWKV与其他Transformer变体的复杂度对比

RWKV是在Linear Transformer和AFT(Attention Free Transformer)基础上迭代来的，具有如下特点：

1. 改造AFT，通过Liner Transformer变换将self-attention复杂度由$O(N^2)降为O(N)$；
2. 保留AFT简单的attention形式和Sequential Decoding，具有RNN表现形式。

更多细节可以参考博客《RWKV解读：在Transformer的时代的新RNN》，在此不再赘述。

进入25年，Transformer还是一统天下的局面，但是各种线性注意力的变种也开始流行起来，1月份minimax也发布了最新的模型提出了Lightning Attention，将支持的上下文长度提升到400万，它能否获得同行的一致认可，我们拭目以待。

5. 其他

5.1 Prompt压缩

对长上下文进行压缩可以极大降低KV Cache的规模。一种可行的方案是首先对长上下文分段利用大模型提取摘要，然后将每段的摘要组合起来再次灌入大模型进行推理，只要摘要没有丢失重要信息这种方案的效果和长上下文接近。这种方案算是一种比较通用的压缩上下文的方法，在AI Agent中经常会采用。

5.2 CPU推理的可能性

当前最好的GPU H100的显存是80G，但是高性能服务器的内存则可以轻易达到1T甚至2T。所以从存储上来说，服务器的内存规模要远大于GPU的显存。理论上，除了上面提到的KV Cache压缩方法之外，我们也可以考虑在CPU上进行大模型推理。一种可行的异构方案是：prefilling阶段还是利用GPU完成，但是后续的decoding阶段则在CPU上实现。利用CPU的内存来存储KV Cache就可以极大缓解存储压力，将batch增大到很大。但这里存在两个难题：1. CPU内存的带宽要远小于GPU显存的带宽，而大模型推理的decoding阶段恰恰又是访存占瓶颈，所以解码速度可能会是个问题；2. 利用CPU进行推理的速度可能要比GPU慢不少，虽然Intel的高端CPU也内置了AMX计算单元，但是算力还是要比GPU差一个数量级。所以即便CPU的大内存可以缓解KV Cache带来的存储压力，但是CPU推理也有其自身的难题。这方面的资料大家可以看一下量子位发的文章《想跑千亿大模型？算力厂商放大招！CPU通用服务器成为新选择》和《拿CPU搞AI推理，谁给你的底气？》。不过我个人觉得CPU+GPU异构计算有可能是大模型推理的一个方向。

事实上，现在的推理框架确实也可以支持将KV Cache放到内存里。HF的Transformer框架就可以支持将KV Cache保存成OffloadedCache。它只将当前正在计算层的KV Cache保留在GPU显存中，其他层的KV Cache都放在内存里，然后通过流式的方式将内存中的KV Cache预取到显存中。这样就可以大幅降低显存的使用，但是速度会慢于KV Cache全在显存中的DynamicCache存储方式。

5.3 vLLM的PagedAttention

在介绍PagedAttention论文时，我们提到vLLM借助PagedAttention可以非常容易实现不同解码策略下的KV Cache共享，这在某种程度上来说也是一种Cache压缩。另外，PagedAttention提升了显存利用率也是变相地压缩了KV Cache，具体看参看《大模型推理–PagedAttention》。

6. 反思

这篇博客本应24年9月份完成，但是入职新公司之后就变得懒惰，遂一拖再拖，拖到现在再看就会发现部分博客的内容已经过时……我个人觉得进入25年之后，大模型的推理技术迭代速度大幅加快，感觉3、4个月就会出现一个巨大的技术突破，如果不能持续跟进就会出现“一步跟不上步步跟不上”的被动局面。不期望自己成为行业大佬，只希望自己还有能力跟上技术的步伐，与大家共勉~

7. 参考

1. ChatGLM2-6B：性能大幅提升，8-32k上下文，推理提速42%
2. 最强英文开源模型Llama2架构与技术细节探秘
3. 大模型百倍推理加速之KV cache篇
4. 聊聊大模型推理中的 KVCache 压缩
5. KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches
6. 一文了解Mamba和选择性状态空间模型 (SSM)
7. Mamba模型讲解
8. RWKV-7-0.4B 模型正式发布，社区发布多款基于 RWKV-7 的新项目！
9. RWKV: Reinventing RNNs for the Transformer Era
10. 如何评价最新的RWKV论文 (arXiv 2305.13048)？
11. An Attention Free Transformer
12. 聊聊大模型推理中 KVCache 压缩方法的性能评测