探秘Transformer系列之（25）--- KV Cache优化之处理长文本序列

万妙音 · 5 天前

探秘Transformer系列之（25）--- KV Cache优化之处理长文本序列

目录

探秘Transformer系列之（25）--- KV Cache优化之处理长文本序列
- 0x00 概述
- 0x01 优化依据
  - 1.1 稀疏性
  - 1.2 重要性
  - 1.3 小结
- 0x02 稀疏化
  - 1.1 分类
  - 1.2 静态稀疏化
    - 1.2.1 Window Attention
    - 1.2.2 StreamingLLM
      - Attention sink
      - 思路
      - 问题
  - 1.3 动态稀疏化
    - 1.3.1 H2O
      - 动机
      - 方案
      - 优化点
    - 1.3.2 keyformer
      - 动机
      - 方案
        
        评分函数
        
        算法
      - 优化点
    - 1.3.3 小结
  - 1.3 针对prefill的稀疏化
    - 1.3.1 FastGen
      - 洞见
      - 算法
    - 1.3.2 SnapKV
      - 洞见
      - 方案
  - 1.4 针对层特点的稀疏化
    - 1.4.1 PyramidKV
      - 洞察
      - 方案
    - 1.4.2 PyramidInfer
      - 动机
      - 方案
    - 1.4.3 ZigZagKV
      - 动机
      - 空间分配
    - 1.4.4 Layer-Condensed KV Cache
  - 1.5 其它方案
    - 1.5.1 针对头特点的稀疏化
    - 1.5.2 投机采样稀疏化
    - 1.5.3 聚类稀疏化
- 0x02 复用
  - 2.1 KV Cache合并
    - 2.1.1 层间合并
    - 2.3 层内合并
  - 2.2 前缀复用
    - 2.2.1 需求
    - 2.2.1 Prompt Cache
      - 洞见
      - 方案
    - 2.2.2 Hydragen
    - 2.2.3 ChunkAttention
    - 2.2.4 AttentionStore
    - 2.2.5 Radix Attention
- 0x03 依据检索的方案
  - - 3.1 InfiniGen
      - 问题
      - 洞察
      - 设计
        
        总体方案
        
        离线阶段
        
        Prefill阶段
        
        解码阶段
        
        KV Cache资源池管理
- 0xFF 参考

0x00 概述

即便做了KV Cache，在长序列场景下，也只是减少了K和V本身结果计算量。但是和后面的部分依然计算量很大，并且K和V存储的压力也非常大。这其实严重阻碍了我们使用更长的序列做输入以及生成更长的序列。要解决这个问题，我们就需要对KV Cache进行优化。
KV Cache优化的核心在于减少内存消耗的原则。这一目标可以通过进一步压缩键值对（Key-Value Pairs）中的“K”（键）或“V”（值）来实现。这些压缩技术的选择直接影响模型的效率，特别是在内存使用和处理速度方面。压缩键值对就是让K和V变短，也就是说我们要舍弃或者压缩一些K和V。通过丢弃这些 token ，我们实际上将相应的注意力分数设置为零，并将注意力矩阵近似为一个更稀疏的矩阵。
减少序列长度是KV-Cache压缩领域研究最多的点。目前大致分为几个大方向：

抛弃不重要的token（稀疏化）。通过attention score或者attention的系数特征来判别不重要的token，然后把这些token对应的KV都抛弃掉。这部分研究跟LLM的量化一样，都在activation outlier提出来之后取得了比较大的进展。虽然这种方法可以提高效率，但存在丢弃关键词元的风险，特别是对于需要深入理解远距离上下文的任务。
复用token。通过复用来降低序列长度。比如复用system prompt等，这是一种工程优化，从算法层面上看，是一种无损优化，无需训练侧介入。
基于检索（retrieval/取回）的方法。比如把KVCache offload到cpu，以页或者聚类的形式组织。q和每个页或者聚类进行相似度计算来决定使用哪些页或者聚类。每次只把得分最高的top-k个页或者聚类加载到显存计算。

注：全部文章列表在这里，估计最终在35篇左右，后续每发一篇文章，会修改此文章列表。
cnblogs 探秘Transformer系列之文章列表
0x01 优化依据

减少序列长度的理论根源主要是以下两点:

大型语言模型（LLMs）注意力机制中的内在稀疏性。
token的重要性程度不同。

1.1 稀疏性

尽管transformers存在巨大的计算需求，但LLMs的注意力机制内部确实存在固有的稀疏性。同理，KV Cache中做为注意力机制的历史记录，也必然是很稀疏的。当然，稀疏性的程度可能因特定的下游任务而异。
下图(a) 展示了在使用CNN/DailyMail数据集进行摘要任务时，不同模型之间注意力稀疏性的多样化水平。这种变化在模型的各个层面都有体现，包括整个模型、单个层以及模型的特定部分。在(b)中，累积分布函数（CDF）描绘了注意力分数与总上下文占比之间的关系。可以看出来，即使KV缓存大小减少了50%，关键token的平均累积注意力得分仍然保持在一个相当高的水平，大约在90%到95%之间。进而也可以产生同样的token。
论文“H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models”的作者也发现，KV Cache其实是很稀疏的，只要5%就可以达到相同效果，即产生同样的token。
因此，我们可以在内存消耗和模型准确性之间做一定的妥协和权衡。

1.2 重要性

导致稀疏性的一个原因可能是因为token的重要性不同。
LLM 最核心的模块是 Attention：通过 Q 、 K 两个矩阵进行矩阵乘，计算$O(N^2)$大小的注意力得分矩阵，并使用注意力得分矩阵对 V 矩阵求加权平均值，得到注意力层的输出矩阵。其中注意力得分矩阵每一行表示 Q 中每个 Token 和 K 中每个 Token 的相关性得分，注意力层输出矩阵的每一行表示每个 Token 与其他 Token 的加权平均值。即然是加权平均，就说明每个 Token 的重要性并不相同，有的 Token 权重更大，而有的 Token 权重更小。
实际上，在文本生成过程中，只有一小部分 Token 获得了最多的注意力。这强调了特定关键token的重要性以及它们在理解上下文和促进文本生成中的核心作用。为了发现这些token，不同研究人员做了不同的工作。

有些工作已经显示了对初始token的偏向性。这种偏向性源于在解码迭代过程中累积的注意力分数倾向于初始token的累积效应。即无论它们与语言建模任务的相关性如何，初始tokens都被分配了惊人的大量注意力得分（很大的原因并非是 token 决定的，而是其起始位置决定的），人们将tokens称为“attention sinks”。当文本长度超过额定长度时，头部的 token 就会被遗弃掉，这就会在 softmax 阶段产生很大的问题，进而影响后续的推理结果。
也有一些工作发现稀疏的KV缓存可以有强时间局部性。稀疏性分布满足幂律分布，少部分固定的token一直主导着decoding。即在计算attention score时，只有一小部分token对模型的价值贡献最大，而且这些token在整段生成中都是重要的。人们称这些token为“Heavy-Hitters”（H2）。比如，论文“Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time”发现，单个 token 的注意力分数呈现强幂律分布，这意味着某些 token 对 Attention 机制的重要性远高于其他 token。作者据此提出了“重要性的持久性”（Persistence of Importance）假设：在生成过程中，只有具有显著重要性的 token 会对后续生成步骤产生持续影响。论文进一步指出，大多数 Transformer 层之间的注意力模式高度相似，重叠比率超过 90%。因此，通过识别“高权重分数 token”，可以有效减少 KV Cache 的内存占用。此外，还可通过预先分配 KV Cache 的空间，仅存储必要的 token，从而将计算复杂度简化为常数级别。

当然，也有另一种可能是：因为模型经过训练，不需要关注整个序列（例如，Mistral AI的Mistral-7B），因此会对某些token不再进行关注或者很少关注。
1.3 小结

我们已经知道，注意力模块倾向于不成比例地持续将更多注意力集中在序列中的少数几个词元上。那么我们就可以不去存储模型几乎不关注或很少关注的词元的键和值。
通过丢弃 KVCache 中一些权重小的 Token，让它们不再参与注意力计算，我们将相应的注意力分数设为零，并用一个更稀疏的注意力矩阵来近似注意力矩阵。这样就可以在保证模型效果的前提下压缩 KVCache 长度，从而在一定量的显存下保存更多的 Token，提高长文本的推理效率。另外，不将序列中所有词元的键和值都存储起来的一个可能的原因是：我们可以在每次迭代中重新计算缺失的键和值，即用时间换空间。
因此，压缩 KVCache 长度的问题就可以转化为通过算法设计找出一个序列中相关性更高的 Token。
0x02 稀疏化

减少序列长度的主要手段就是稀疏化。KV Cache 稀疏化（KVCache Sparsification）是通过减少不必要的存储和计算来提高效率的一种优化方法。它的核心思想是选取需要关注的关键 token，并丢弃无关的 token，这是一种算法层面的优化，直接降低了计算，以减少 KV Cache 的体积和计算复杂度。
1.1 分类

KV Cache 稀疏化的各种方案本质上就是选取需要关注的 token 方法的不同。但由于舍弃了部分 Token 的注意力计算，所以理论上是有损的，可能需要训练侧保证模型效果。稀疏化属于目前比较主流的KVCache研究方向，不同论文对其的具体分类也不尽相同。两篇综述论文将稀疏化称为Token Dropping和KV Cache Selection。下图就是两篇综述中对稀疏化的不同称呼和细分。

Token丢弃是一种识别不重要token并将其丢弃的技术。然而，这些方法的关键挑战在于如何确定哪些token是不重要的。
KV Cache选择机制旨在降低KV缓存的内存利用率，最大限度地减少推理延迟，并提高大型语言模型的整体吞吐量。KV Cache选择机制以运行过程进行分类：静态KV Cache选择机制仅在预填充阶段执行token过滤，所选token在后续解码步骤中保持不变；动态KV Cache选择机制在解码阶段不断更新KV缓存，实现自适应缓存管理。

本文对稀疏化主要分为以下：

静态稀疏化。在计算预测下一个 Token 时，只维护一个窗口大小的历史 Token 信息。然后通过调整注意力窗口来调整序列长度。静态稀疏化的窗口内的 Token 是固定的。
动态稀疏化。通过调整注意力窗口来调整序列长度。动态稀疏化窗口内的 Token 是不固定的。
针对prefill阶段的稀疏化。主要是优化prompt的序列长度，取出其中重要的token。
针对层特点的稀疏化。LLM模型的注意力层有自己的特点，这也影响了KV Cache稀疏化策略。因此，在KV Cache管理中，即需要依据注意力权重的贡献进行区分处理或者累积处理，也需要动态调整每层参与注意力计算的键token数量。
其它方案。比如结合投机采样进行的稀疏化；针对头特点的稀疏化等。

我们也用论文”A Survey on Large Language Model Acceleration based on KV Cache Management“来进行比对学习。该论文中与稀疏化对应的部分是KV Cache Selection、KV Cache Budget Allocation和KV Cache Merging。因为与笔者的分类方式不是正交的，所以我们在图上做出标注。

我们接下来通过案例进行学习。
1.2 静态稀疏化

静态稀疏化可以通过注意力机制来实现，即稀疏注意力机制或者线性注意力。稀疏注意力机制的核心思想就是在推理中选择合适的 Token 来进行相应的计算，这种方案在序列比较长时尤其有帮助，可以大幅降低 Attention 部分的 KV Cache 大小和计算量。线性注意力机制（如Linear Transformer、RWKV和Mamba等）通过将标准注意力机制替换为与序列长度线性相关的机制来减少内存需求。然而，这种方法可能会降低模型的表达能力，导致在需要复杂、长距离词元依赖关系的任务中性能下降。
比如，Mistral-7B使用被称为滑动窗口注意力（SWA）或局部注意力的注意力机制变体。局部注意力保证我们不需要关注整个序列，只通过关注最后（4096个）相邻的词元来构建词元表示，这样在KV缓存中永远不会存储超过窗口大小（例如4096）的张量对。再比如，
论文"Efficient Streaming Language Models with Attention Sinks"中分析了四种 Attention 实现，具体如下图所示。其中T 为待预测的第 T 个 Token，L 表示 pretrain 时的最大序列长度（一般为 2k、4k 等），由于主要针对的是长序列场景，因此 T 远大于 L。

第一种是原始注意力实现，后三种是稀疏化实现，这些稀疏化方法都是静态的，即 Token 之间的相关性是一种固定的范式，每个 Token 的相关的 Token 都是固定距离的。

Dense Attention是vanilla Attention 实现。每个 Token 都能看到该 Token 及之前的 Token，其计算复杂度为 $O(T^2)$，KVCache 存储复杂度为$O(T)$。由于复杂度比较高，T 在预训练的时候会比较小。在推理的时候，当文本长度超过了预训练时的长度，模型效果就会大幅降级，所以表现出来的 PPL 值也比较大。
窗口注意力（Window Attention）缓存最近的 n 个tokens的KV。这种方法相当于只在最近的tokens的KV状态上维护一个固定大小的滑动窗口。其中灰色的虚线框表示超过窗口长度后从 cache 中淘汰的 Token。计算复杂度为 O(T*L)，计算复杂度低，cache 比较小，在推理中效率高。但是模型效果很差（PPL 很大），因为虽然但一旦开始tokens的键和值被删除，性能就会急剧下降。
带重计算的滑动窗口（Sliding Window with Re-computation）。这种 Attention 与 Window Attention 类似，区别是 Sliding Window 不缓存窗口的 Tokens，而是为每个新token重建来自 $L$个最近tokens的$KV$状态。这种方案的模型效果对比Window Attention 高，PPL 值远低于 Window Attention。主要原因是重计算把窗口中的 Tokens 作为初始 Tokens，这样既保留初始 Tokens 又保证计算只在一个窗口内，大大降低了KVCache 存储复杂度。由于重计算的存在，其精度可以保证，但是性能损失比较大。虽然它在长文本上表现良好，但由于在上下文重计算中的二次注意力，其的$O(TL^2 ) $ 复杂度使其相当缓慢，使得这种方法不适用于实际的流式应用。
StreamingLLM 保留了用于稳定注意力计算的attention sink（几个初始tokens），并结合了最近的tokens。它高效并且在扩展文本上提供稳定的性能。黄色的框表示初始的几个 Token，也就是 Attention Sink，其计算复杂度为 O(T*L)，模型精度很不错，和带重计算的滑动窗口 Attention 相当。

我们接下来分析窗口注意力和StreamingLLM。
1.2.1 Window Attention

Window Attention采用滑动窗口机制来解决长文本推理挑战，其中落在窗口外的token被永久驱逐并变得无法访问。
在原始注意力中，每一个token，都要和它之前所有的token做Attention。而通过我们平时语言习惯中可以知道，一段话中每个字之间的相关性差别很大。对于当前token，一般来说越相近的字相关性越强，距离越远的token，能提供的信息量往往越低，因此似乎没有必要浪费资源和这些远距离的token做Attention。因此，滑动窗口（Sliding Window Attention）机制就是：每一个token只和包含其本身在内的前 L（L 表示窗口长度）个token做Attention，即只和当前token距离相近的token做注意力计算。这样就可以把cache的容控制在 L 内。因为每个 Token 只和邻近的 Token 做 Attention 计算，所以计算复杂度为 O(TL) ，KVCache 存储复杂度为O(L)。这种方法极大的降低了 KVCache 的存储开销，从线性复杂度降低到常数复杂度。
典型代表即 Longformer，具体如下图所示。

(a)展示了传统的“全注意力”机制，其中每个新生成的token都会关注序列中所有先前的token。
(b)描述了“滑动窗口注意力”。该方法维护一个固定大小的最近token滑动窗口，给定固定的窗口大小w，每个token都会关注其两侧$ \frac{1}{2}w $个token。滑动窗口允许我们拥有一个固定大小的缓存。一旦序列增长超出滑动窗口token数，我们可以在缓存中循环并开始覆盖。由于w可以设置的比较小，所以可以和输入序列长度n呈线性关系，其计算复杂性为$O(n×w)$，更加适合浅层捕获局部信息。而且，因为缓存中的位置无关紧要，所有与位置相关的信息都通过位置嵌入编码，所以很容易实现并且效果很好。然而，这种方法限制了模型从过去捕获全面语义信息的能力，导致文本生成质量较低且准确性下降。
人们可能会疑惑，虽然距离越远的token涵盖的信息量可能越少，但不意味着它们对当前token一点用处都没有，在滑动窗口中直接杜绝了它们的参与，这真的合理吗？其实，Silding Window Attention并非完全不利用窗口外的token信息，而是随着模型层数的增加，间接性地利用起窗口外的tokens。这里可以类比为CNN技术中的感受野。如果堆叠多个层，从layer0开始，每往上走一层，对应token的感受野就往前拓宽W，虽然在每一层它“最远”只能看到前置序列中部分token，但是只要模型够深，它一定能够在某一层看到所有的前置tokens。该层窗口注意力可以访问所有输入位置，产生一个大的感受野，可以构建跨整个输入信息的表示。在具有l层的Transformer中，顶层的感受野大小为$l \times w$(假设对于所有层w是固定的)。根据应用程序的不同，可能为每个层使用不同的w 值更好，以在效率和模型表示能力之间取得平衡。具体如下图所示。
(c)展示了一种称为“膨胀窗口注意力”的变体，其准确性与窗口注意力类似，也存在类似的限制。其特点如下。
- 为了进一步增加感受野而不增加计算量，滑动窗口可以"扩张"。这类似于dilated conv（空洞卷积或扩张卷积），其中窗口具有大小为扩张值d的间隔，这样可以将感受位置的间隔放大，让token在当前位置适当捕捉到得更远处的信息，关注更大的视野，考虑了更加全面的上下文信息。由于考虑了更加全面的上下文信息，空洞滑窗机制比普通的滑窗机制表现更佳。
- 每个位置的Q关注的K还是W个，但是W个不是连续的而是跳跃的。定义跳跃的间隔是d, 那么Q关注的K的范围就是w * d。计算方面依然保持滑动窗口在计算性能方面的优点。
- 假设对于所有层都是固定的d和w，那么感受野大小为$l × d × w $，即使对于较小的d值，也可以触及成千上万个token。在多头注意力中，每个注意力头计算不同的注意力分数。每个头部具有不同的扩张配置的设置可以通过允许一些没有扩张的头部关注局部上下文，而其他具有扩张的头部专注于更长的上下文，从而提高性能。
(d)是Global+sliding window（融合全局信息的滑窗机制），滑动窗口和扩张注意力不足以学习特定任务的表示，需要从其它方面借鉴。因此引入了全局注意力对这两者进行补充。其特点如下：
- 全局注意力就和普通的transformer是一样的，都是关注全部的K。依据具体任务的不同，在特定选择的位置引入关于全局的计算，使得模型在计算过程中能够接收到全局的信息。
- 具有全局注意力的token会关注整个序列中的所有token，而序列中的所有token也会关注它。由于这些token的数量相对较小且独立于n，因此结合局部和全局注意力的复杂度仍然为O(n)。虽然指定全局注意力是与任务相关的，但这是一种向模型的注意力添加归纳偏差的简单方式，比起使用复杂架构将信息跨越较小的输入块组合的现有任务特定方法要简单得多。

1.2.2 StreamingLLM

StreamingLLM 利用“注意力沉积”效应，用早期Token 的KV结合近期上下文来优化长序列处理，实现了对无限长度输入的支持，同时生成无限长度的输出。
StreamingLLM 来自论文《Efficient Streaming Language Models with Attention Sinks》，其目标是让模型能够处理无限长度的输入（这里的“无限长度输入”与“无限长度上下文”有所不同——前者无需记住所有输入内容）。
Window Attention方案虽然压缩了KVCache的长度被压缩，但是模型效果却不好，主要原因是最初始的 Token 被丢弃了。这种 Token 的重要性其实非常高，丢弃会严重影响模型效果，精度下降比较大。而StreamingLLM恰恰弥补了这个缺憾。StreamingLLM发现注意力中的一个关键现象，即初始序列token中保留的键值对保持了关键的模型性能。这种注意力下沉效应通过早期位置的不对称注意力权重积累来体现，而与语义意义无关。该方法利用了这一特性，将注意力汇位置与最近的上下文相结合，以实现高效处理。
StreamingLLM策略在 Window Attention 的基础上，通过仅保留最初的位置词元（sink tokens）和最后相邻的词元（局部注意力）来构建了一个滑动窗口缓存。因此，StreamingLLM KV Cache 的长度是固定的，既有固定部分（通常1到4个 token ），也有滑动部分。这样既保留 Window Attention的特性，使得 KVCache 存储复杂度为O(L)，计算复杂度为 O(TL) ，同时又保证了模型效果不因丢失初始 Tokens 而大幅下降，PPL 和 Sliding Window Attention with recomputation 相似。
Attention sink

StreamingLLM作者发现，window attention只要加上开头的几个token比如就4个，再操作一下位置编码，模型输出的PPL就不会变好，输出到20k长度都很平稳。论文将这种现象叫做Attention Sink（注意力池）。就好像文本的注意力被“沉溺”到了开头的位置。即初始词元聚集了大量注意力。而且intial tokens与生成token的绝对距离距离和语义信息都不重要，重要的是intial tokens是第一个或者前面几个token，就是这些token作为锚点，所以其权重特别大。

上图给出了lama-2-7B中256个句子的平均注意力逻辑的可视化，每个句子的长度为16。观察结果包括：

前两层（第0层和第1层）的注意力图呈现出“局部”模式，最近的token受到了更多的关注。
除了前两层之外，之后的注意力主要集中在first token。

如果删除这些初始token的KV，将导致注意力计算中SoftMax函数的相当一部分分母也被删除。这种变化导致注意力得分的分布发生了显著变化，偏离了正常推理环境中的预期。
为何会出现这种现象？Evan Miller 在“Attention Is Off By One” 做了精彩的解读。

The problem with using softmax is that it forces each attention head to make an annotation, even if it has no information to add to the output vector. Using softmax to choose among discrete alternatives is great; using it for optional annotation (i.e. as input into addition) is, like, not cool, man. The problem here is exacerbated with multi-head attention, as a specialized head is more likely to want to “pass” than a general-purpose one. These attention heads are needlessly noisy, a deafening democracy where abstention is disallowed.

在Attention机制中，Softmax的输出代表了key/query的匹配程度的概率，如果softmax在某个位置的值非常大，那么在反向传播时，这个位置的权重就会被大幅度地更新。然而，有时候attention机制并不能确定哪个位置更值得关注，但是Softmax又需要所有位置的值的总和为1，因此模型必须“表态”给某些位置较大的权重，因此，模型倾向于将不必要的注意力值转嫁给特定的token，这些token就是initial tokens。
因此，Evan Miller改进了一下Softmax，提出了“SoftMax-off-by-One”：把softmax的分母加了个1，这样所有位置值可以加和不为1，Attention就有了可以不对任何位置“表态”的权利。

思路

根据attention sinks特性，论文作者参考之前多轮对话场景的解决方法，给出StreamingLLM的解决方案，即在window attention基础上，新引入了initial tokens的KV。

锚点加窗口。StreamingLLM中的KV cache可以分为两部分，既有一个固定部分（通常为1到4个词元），又有一个滑动部分。
- 固定部分。保留整个序列的初始 Tokens（attention sink部分），图中画的是4个initial tokens；其作用是识别并保存模型固有的attention sink，锚定其推理的初始token来进行注意力计算并稳定模型性能。
- 滑动部分。滑动窗口的KV（Rolling KV部分）cache，保留了最近的3个tokens，窗口值固定为3。其目的是：想让大模型无限输出，达到“Streaming”效果：不用KV cache会计算太慢，但是用了KV cache，显存占用随长度增长太多。因此需要丢弃一些KV cache。
计算。每个 token 只和窗口内的 Tokens 以及序列的初始 Tokens 进行 Attention 计算。而且训练期间，用“SoftMax-off-by-One”替代常规softmax，解决attention sink问题。

在实现中，会用“Rolling Buffer Cache"技术来丢弃KV cache中的某些行，同时也会控制位置编码。下图给出了Rolling Buffer Cache的运作流程，当我们使用滑动窗口后，KV Cache就不需要保存所有tokens的KV信息了，你可以将其视为一个固定容量（W）的cache，随着token index增加，我们来“滚动更新” KV Cache。

问题

StreamingLLM虽然记住了Attention sink以及最近的窗口里面的token对应的KV cache，然而仅仅在注意力机制中依赖这k个token并不能提供必要的准确性。因为目前是按照位置把中间的token都丢掉了，万一中间的token就是重要的怎么办。这就是静态 Token 稀疏化的问题：Token 候选集过于固定。这种 Token 候选集的设计源于作者通过观察某些数据集中 Token 之间的相关性发现的规律。这种规律不能保证具备普适性，会导致关键token注意力中近期上下文的丢失和窗口注意力中关键上下文的缺失，会损害模型在长序列任务上的性能。
1.3 动态稀疏化

动态 Token 稀疏化本质是为当前处理 Token 维护一个相关性高的历史 Token 集合，但与静态 Token 稀疏化不同，这个集合的构造不再由 Token 距离或者固定 Token 决定，而是设计一种算法去筛选历史的 Token。静态 Token 稀疏化是动态 Token 稀疏化的子集，所以理论上动态 Token 稀疏化的效果会比静态 Token 稀疏化更好。
然而，在推理过程中，尤其是当输入序列包含未知或未见过的token时，动态确定哪些token作为关键token是一个相当大的挑战，针对这挑战主要有两种方法：

评分函数。利用评分函数识别关键token。比如我们为每个token引入了一个评分函数fθ，以从总共n个token中识别出k个关键token。
驱逐策略。保留系统认为重要的KV，舍弃那些不重要的。其中重要性的判断非常重要，大部分的系统都是对于每一层按照attention score的大小，保留top-B的结果。

其实这两者是一枚硬币的两面：通过动态的评价策略来判断需要保留和废弃的KV值，通过运行时键/值驱逐来减少KV缓存大小。

需要注意的是，大部分工作都假设了注意力模式在不同迭代之间的持久性，即，如果某个token在当前迭代中被认为是不重要的（即具有较低的注意力权重），那么在生成未来token时，它也很可能保持不重要。基于这一假设，当KV缓存大小超过其预算时，他们会在每次迭代中从KV缓存中驱逐具有低注意力权重的token。被驱逐token的键和值在从内存中删除后，将永久排除在后续迭代之外。
我们通过几个案例来学习下。
1.3.1 H2O

论文”H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models“提出的H2O是比较经典的方案。H2O观察到，注意力计算主要由一组被称为重拳（“pivotal tokens”或“heavy hitters”）的高影响力token来驱动。因此，H2O将缓存优化重新表述为动态子模型问题，利用累积注意力得分来指导token保留决策。
动机

稀疏的KV缓存并不一定具有100%的时间局部性。比如，LLM在decode第6个token时不一定需要第1至第5个token，而是有可能需要其中的任意数量的token。想象一种最极端的情况，decode第6个token时，需要第1、2、4个token，decode第7个token时需要第3、5、6个token。因为在同一时刻虽然不会使用所有token，但是每个token都有可能在任意时刻被使用，所以这种情况下必须在内存中保存所有历史token的KV缓存。论文将这类情况称为Dynamic Sparsity，见下图左一。
一种直观的想法是，也许稀疏的KV缓存可以有类似密集KV缓存那样的强时间局部性：在两个decoding步骤之间仅有一小部分KV缓存不同，并且在某个decoding步骤被稀疏掉的token永远不会在接下来再被使用，这样就实现了与密集KV缓存类似的时间局部性同时减少了KV缓存内存占用。论文将这类情况称为本地稀疏（Static Sparsity (Local)）和等步长稀疏（Static Sparsity (Strided)），见下图左二左三。
H2O认为，本地稀疏、等步长稀疏等强时间局部性的稀疏都太naïve了，光看最近的K个元素或者等步长最近的K个元素无法代表所有的历史信息。在满足强时间局部性的稀疏策略中，应该存在更优的策略。H2O利用累积的token重要性在decoding过程中找到了这一策略，见图左四。

H2O作者认为，最为理想的drop token方式是：预知未来的token对当前token的注意力分数，然后drop掉那些分数低的token。因为未来的token还不知道呢，所以只能预测。
因为attention分数比较大的token对计算结果的影响比较大，所以作者拿前面的attention分数相加的结果来预测之后的attention分数。作者在论文后续附录中证明了在submodular情形下，H2O假设的贪心地构造KV缓存的策略是near-optimal的。即随着规模的增长存在边际效应。因此，H2O的设计问题已经收敛到了一个很小的问题：dynamic submodular problem（动态次模问题）。即如何设计贪心算法，满足每次驱逐的token都是submodular最优的。即每次丢掉的token对生成的影响最小。
方案

H2O的设计目标就是确定一个KV缓存维护/驱逐策略。H2O设置最大缓存 token 数量（预算），每次达到缓存预算时，该策略可以通过贪心法去计算token的打分，每次丢弃累积注意力分数最低的 token ，因此只保留在迭代中始终实现高注意力分数的 token。这样使基于稀疏KV缓存的decoding效果与基于密集KV缓存的decoding效果相近。
H2O作者作者发现：在给定步骤中具有影响力的词元（“pivotal tokens”或“heavy hitters”）在未来步骤中仍然具有影响力。换句话说，我们可以确信被丢弃的低影响力词元在未来生成步骤中仍然可以相对被忽略，因此可以安全丢弃。基于此观察，作者提出了一种名为“Heavy-Hitters Oracle”（H2O）的贪心的历史 Token驱逐策略：对于某个token A ，A以及A之前的token都会对A有注意力，把这些注意力加起来，可以得到A的“重要性分数”。将每个token的重要性分数都计算出来，然后丢弃分数最低的那些token即可。H2O选择了注意力分数作为评分函数，记作$f_θ(acc; attn)$。这样就可以识别出在提示和token生成阶段持续获得较高注意力分数的token作为最关键或关键的token。此处其实有提高空间，因为是使用attention score的绝对值来反映token的影响力，求和隐含了“越久的token越重要“的假设，越久的token求和项越多，最终的值有可能越大。
算法具体流程如下图所示：

初始化相关性历史 Token 集合，设置集合最大容量。
迭代更新历史 Token 集合。
- 当历史 Token 集合大小等于最大容量时，将当前 Token 加入历史 Token 集合中；
- 当历史 Token 集合大小大于等于最大容量时，计算当前 Token 和历史 Token 集合中所有 Token 的相关性得分（实现起来是对于注意力矩阵按列求和再取topk），将最低分的 Token 从集合中剔除，将当前 Token 加入到集合中。即每生成一步会将当前 Token 相关性最低的历史 Token 永久剔除，被剔除的历史 Token 将不会参与后续任何一个 Token 的 Attention 计算。用自然语言描述，就是KV缓存中被稀疏掉的部分对应的attention score设置为0，其它部分正常计算。
这样可以保证前后两个KV缓存的大小相等，最多有一个位置不同。即：要么从前一个时刻的KV缓存中删掉一个，然后加上本次生成的token的KV缓存；要么保留前一个时刻的KV缓存，不加本次生成的token的KV缓存。这样的KV缓存始终满足强时间局部性。

图上的反斜杠"\"是一种集合运算符号：相对差集。”U \ A“表示在集合U中，但不在集合A中的所有元素，比如相对差集{1，2，3} \ {2，3，4} 为{1} ，而相对差集{2，3，4} \ {1，2，3} 为{4} 。

与H2O类似的方案是Scissorhands。H2O算法一次仅丢弃一个词元，而Scissorhands则根据目标压缩比例（例如，减少30%的KV缓存）丢弃尽可能多的词元。Scissorhands简单地保留最近的词元和在历史窗口内具有最高注意力分数的词元。而H2O则丢弃累计注意力分数最低的词元，只保留那些在迭代过程中始终具有高注意力分数的词元。
优化点

现有通过驱逐token来减少KV Cache大小的研究本质上存在一些挑战。类似H2O的方法假设注意力模式在不同迭代之间不会改变，换句话说，H2O其实是用上一层KV缓存的稀疏性来指导下一层的KV缓存的稀疏性。但实际上情况并非如此，因为注意力模式在迭代之间具有动态性，在当前迭代中被认为不重要的token在后续迭代中可能会变得重要。在前面生成过程中被剔除的 Token 可能与当前的 Token 有着较高的相关性，但因为 Token 在前面的生成过程中已经被剔除，导致无法把 Token 召回和当前 Token 完成 Attention 计算。
因此，H2O的评估窗口较窄，导致其在KV缓存预算范围内表现出较高的效率，但随着序列长度超出KV缓存预算，它开始难以应对注意力模式的动态性，未能很好捕捉这些动态变化。
另外，H2O方法将保留的键/值token数量设置为输入序列长度的一个固定百分比。在这种情况下，无论生成了多少个token，KV Cache预算始终保持恒定。然而，这种固定KV Cache预算存在一些局限性。每一层所需的键/值token数量是不同的，而每个查询token也需要不同数量的键/值token来有效地表示基线模型的注意力模式。此外，即使是相邻的查询token，所需的键token数量也会发生显著变化。固定的KV Cache预算忽视了查询token之间的这种差异性，如果保留的键/值token数量不足以准确表示基线模型的注意力权重，会导致KV Cache管理效率低下。因此，需要根据每个查询token的需求，动态调整所加载和计算的键/值token数量，同时考虑不同层和查询token之间的差异性，以实现更高效的KV Cache管理。
1.3.2 keyformer

在生成推理中，大约90%的注意力权重集中在特定的一小部分token子集上，这些token被称为“关键”token。这些token对于大型语言模型（LLMs）理解上下文至关重要，但可能不在窗口注意力的滑动窗口内。Keyformer引入了一种混合注意力方法，如下图(d)所示，该方法在生成下一个token时结合了最近的token和前面的关键token。此外，Keyformer揭示了token移除会扭曲底层softmax概率分布。考虑到softmax分布在token显著性评估中的关键作用，Keyformer 结合了正则化技术来减轻这些分布扰动。

动机

下图比较了完整注意力下的得分函数分布与KV缓存减少后的分布情况。如方程2所示，当KV缓存减少时，分数较低的token会被丢弃，因为isoftmax函数的性质，这会破坏得分函数的分布，进而导致上下文信息的丢失、准确性的降低以及可能生成的文本质量下降。因此需要增强模型提取近期关键token的能力。

方案

评分函数

Keyformer通过对未归一化的对数几率（logits）应用正则化来调整由于丢弃tokens而导致的评分函数变化，从而识别关键tokens。具体而言，Keyformer提出了一个新颖的评分函数，记作fθ(Keyformer)，以解决基于累积注意力的评分函数（fθ(acc attn)）的局限性。这个新的评分函数将Gumbel噪声分布整合到未归一化的logits中。然而，它在形成基础概率分布时未能考虑被丢弃的token。为了纠正这一点，Keyformer引入了一个温度参数，记作τ。
具体剖析如下：

Keyformer使用对数几率正则化技术。通过引入额外的分布（ζ）来正则化减少的对数几率，使模型保持鲁棒性和适应性。即使在推理过程中存在未知上下文的情况下，它也能帮助识别关键token。Keyformer将这种噪声添加到从QKT派生的未归一化对数几率中。此外，添加的分布类型对最终的概率分布产生显著影响。对应下图标号1。在这里，yi是经过调整的对数几率（logits），xi是未归一化的对数几率，ζi是用于正则化的额外分布。
计算概率时，用温度参数 τ 考虑被丢弃的token。“温度”参数（τ）在调节概率分布的平滑性方面起着关键作用。这个参数控制概率中的随机程度。在token从KV缓存中被移除时，τ 尤为重要，因为它们不能在没有重新计算其键和值的情况下被重新引入。较高的τ值（τ → ∞）产生均匀概率，将所有token赋予相同的分数。相反，较低的τ值（τ → 0）产生更尖锐的分布，根据未归一化的logits优先考虑特定的token。另外，随着更多的token被丢弃，我们需要一个更均匀或随机的概率分布。因此通过逐步增加τ和∆τ来系统地增加我们评分函数fθ中的随机性。
选择的分布受到Gumbel分布的启发。Gumbel分布特别适合我们的关键token识别任务，因为它描述了样本集中最大值的分布，并且偏向于初始token。这使得它成为长序列中建模关键token的理想选择。标号3展示了应用于未归一化logits的标准Gumbel概率密度函数（pdf），而标号4显示了添加Gumbel噪声后logits的pdf。

算法

尽管当前token的相关性在识别关键token方面很重要，但模型在大多数生成的token中应该保持一致的行为。为了基于这种一致的行为来识别关键token，Keyformer在提示阶段和token生成阶段都应用累积评分函数（fθ）。如果没有累积，token只能依赖于当前token与先前token的相关性。
在提示处理阶段，Keyformer计算提示长度$S_n$内所有n个token的键和值，以预测第一个token。给定KV缓存预算，Keyformer保留一个最近的w个token的窗口，从n-w个token的窗口中丢弃n-k个token，并且基于Keyformer评分函数来识别出k-w个token。关键token（k-w）和最近token（w）的组合形成了缩减的KV缓存。
在token生成阶段，Keyformer使用缩减的KV缓存进行操作。第一个生成的token仅关注KV缓存中的k个token，如下图图中的解码步骤2所示。最近窗口w向右移动一个token，而评分函数fθ则累积了来自前一个解码步骤的评分函数。在token生成阶段的每个解码步骤中，模型会从大小为k+1-w的窗口中识别出k-w个关键token。因此，最近窗口中添加了一个token并删除了一个token。
此外，温度参数τ通过增加∆τ来调整评分函数概率分布中移除token的数量。

Keyformer的详细算法参见下图。

优化点

与H2O类似的技术如下：

SubGen是一种为 KV cache开发的高效压缩技术。经验证据表明，attention模块中的关键嵌入存在显著的聚类趋势。基于这一关键见解，SubGen设计了一种具有亚线性复杂度的新颖缓存方法，对关键token采用在线聚类并对值进行在线ℓ (2)采样。
LESS（Low-rank Embedding Sidekick with Sparse policy）会学习原始注意力输出和稀疏策略近似的注意力输出之间的残差，通过将稀疏策略丢弃的信息累积到恒定大小的低等级缓存或状态中来实现此目的，从而允许查询仍然访问信息以恢复注意力图中先前省略的区域。

上述基于永久驱逐的方法面临两个重大限制。

首先，对token的不可逆驱逐可能会损害模型在长序列任务上的性能，特别是在needle-in-a-haystack场景中，这些方法被证明难以适应多回合对话环境。
其次，在解码阶段选择KV缓存会引入计算开销，对解码延迟产生不利影响，并影响端到端加速。

为了应对这些挑战，几项研究侧重于开发解码阶段KV缓存选择策略，而无需永久驱逐。比如：

InfLLM：块级存储全KV Cache，GPU保留关键Token，次要部分卸载至CPU。
SqueezedAttention：离线K均值聚类生成语义键簇，推理时加载相关键。
SparQ：基于查询向量显著元素选择性检索缓存键的隐藏维度，近似计算注意力。

这些方法通常采用多层缓存系统（例如CPU-GPU分层缓存），并利用先进的数据结构和系统级增强功能来优化检索效率，从而在减少GPU KV缓存占用的情况下实现高效推理。为了加速关键token的检索，这几项研究提出了基于索引的方法，以块或集群粒度组织和访问KV缓存，采用多级缓存（如CPU-GPU层次存储）避免永久移除，优化检索效率，从而实现高效的查询和提取操作。
1.3.3 小结

论文“A Survey on Large Language Model Acceleration based on KV Cache Management”给出了静态稀疏化和动态稀疏化方法的一些总结。

1.3 针对prefill的稀疏化

现有的KV缓存压缩方法主要关注生成阶段的优化，忽略了输入阶段KV缓存的压缩。而实践中，大模型应用如对话（特别是RAG）和文档处理的特点是：输入很长，而输出相对较短。输入可能就已经把显存撑爆。或者即便显存还可以，每一步计算也会因为需要交互的token太多而非常慢。因此输入阶段的KV缓存是内存和计算瓶颈所在。比如我们输入的prompt有16K，那么存储这16K文本对应的所有KV，无论对于显存还是计算都有极大的压力。因此我们可以在prefill阶段对KV Cache进行处理，drop一些不重要的token，即压缩prompt生成的KV。而解码阶段就照常进行，这样不需要在解码阶段进行额外的计算，又可以加速。
1.3.1 FastGen

FastGen 通过识别五种注意力模式（局部、标点导向、稀疏分布等），再在预填充期间识别重要的token。针对每个token，FastGen 对不同注意力的特点进行针对性的分析（profiling），从而来辨别注意模块的内在结构。然后，相应地在已识别的不同结构上进行了不同的压缩策略，以自适应的方式构建KV缓存。FastGen不设置缓存预算，而是设置注意力矩阵的最大近似误差，因此专注于模型精度保持。
洞见

以前许多压缩KV缓存的研究都没有利用LLM中复杂的注意力结构。而LLM中的注意力结构有如下特点：

注意模块中有丰富的结构。
不同的注意头通常具有不同的功能，遵循不同模式。
并非所有注意模块都需要关注所有表征。

上述特点表明需要根据每个单独的注意头定制压缩策略。下图左侧是四种常见的注意力结构，右侧是同一层的三个注意头的注意图。其中，特殊词元（绿色）+ 标点词元（橙色）+ 局部注意力（蓝色）。灰色为其它词元。

算法

除了传统的全KV缓存策略外，FastGen考虑了四种基本的KV缓存压缩策略。四种KV缓存压缩策略是：

特殊token（Special Tokens）。在KV缓存中只保留特殊的token，如句子开头token 、指令token[INST]等。
标点符号（Punctuation）。在KV缓存中只保留标点符号，如“.”、“：”、“？”。
局部性（Locality）。此策略驱逐长期上下文，保留最后相邻词元。有些注意力模块主要关注局部上下文，也就是我们常说的 Locality。在此策略中，一旦上下文token和当前token之间的相对距离超过阈值，上下文token的KV缓存将被逐出。阈值由预定义的局部上下文长度预算与输入序列长度的比率所确定。
频率（Frequency）。监控每个token的累计注意力分数，然后将这些分数视为token使用频率，并只将最频繁的token保存在KV缓存中。

FastGen可以分为两步：

首先，在预填充阶段结束时，对模型的注意力层进行剖析（profiling），识别各种注意头的行为和结构模式，在这个profiling的指导下，自适应地为每个头选择最合适的满足误差目标的压缩策略，构建KV Cache。与其他方法一样，FastGen假设已识别的注意力模式将在未来的生成步骤中保持不变。如果误差目标过于严格，无法达到，则FastGen会退回到常规的KV缓存。
然后，在每个生成步骤，不是不加区分地为每个新生成token添加新KV向量，而是根据每个token选择的压缩策略来管理其KV缓存。比如驱逐注意头上强调局部上下文的长程文本，丢弃以特殊token为中心的注意头上的非特殊token，并且仅对广泛关注所有token的注意头使用标准KV缓存。

1.3.2 SnapKV

SnapKV 的核心思想是根据用户问题识别重要的token。LLMs在生成过程中对输入tokens的关注模式具有很强的一致性和稳定性，而这种模式可以在输入序列末尾的一个小窗口内就观察到。因此，SnapKV只保留那些被注意力heads持续关注的输入tokens对应的KV，从而大幅压缩KV缓存的尺寸。同时，SnapKV还引入了池化等聚类技术，既能捕捉局部的完整信息,又能过滤掉大部分无关tokens。
洞见

论文作者做了实验，把输入按照128个token为单位来划分window，取靠近尾部的20个window来做计算。对于prefix里面的任意一个token，假设叫做$t_i$。我们取20个window里面的某一个window，假设叫$w_i$。$w_i$的每一个元素和$t_i$都会有交互，因为$t_i$在$w_i$的元素前面。那么，$w_i$的每个元素和$t_i$都会计算出一个attention score。对于窗口$w_i$，$t_i$会有128个attention score。取这128个attention score的均值作为$t_i$的重要性值。这样，对于窗口$w_i$，prefix的每个token都可以计算出重要性值，然后根据平均score选出生成阶段prefix中最重要的topK个token。
然后，作者比较窗口选择出来的重要token和生成阶段选择出来的重要token的重叠率。通过比较20个窗口选择的重要token和生成阶段实际选择出来的重要token，作者发现最后一个窗口选择出来的和生成阶段选择的是高度重合的。因此这就说明：输入里面哪些重要可以使用输入prompt本身就可以决定。即论文所说的：

Pattern can be identified before generation.
Pattern is consistent during generation.

另外，作者还观察到几个现象。

无论生成上下文的长度如何，提示中的特定 Token 在生成过程中始终展现出更高的注意力权重，而且比较稳定。
提示中问题的位置并不会显著改变观察到的注意力模式的一致性。也就是，无论问题的位置如何，都可以获得相关特征的注意力。
注意力模式高度依赖上下文，表明这些模式与用户提出的具体指令有很强的关联。

因此，一个上下文感知的 KV Cache 压缩方法可能会带来更好的性能。
方案

SnapKV简化了FastGen的方法，只专注于根据token的重要性得分检索token。在所有提示token中，只有一部分携带了响应生成的关键信息，这些token在生成阶段保持了其重要性。SnapKV采用末端定位的观察窗口（end-positioned observation window）来检测这些重要的上下文token。然后，将它们对应的键值对与观察窗口中的token连接起来
SnapKV在 Prefill 阶段不是保留所有输入 Token 的 KV Cache，而是采用稀疏化的方式。

在prefill阶段，针对每个 Attention Head 将 Prompt 分为 Prefix 和 Window 两部分，即下图中橙色和绿色部分，可设为representative token 和 window token。对应下图标号1。可以从位于 Prompt 末尾的 “Observation” 窗口获得每个 Attention Head 在生成过程中始终关注特定的注意力特征。
通过 Window 中 Token 与 Prefix 中 Token 的 Attention Score 来选择稀疏化的 Token。即假设window token之前的某个token为A，window token会对A有一个注意力分数，将所有window token对A的注意力分数加起来，就是A的“重要性分数”。重要性分数最高的那些token就是representative token。对应下图标号2。
依据注意力分数选出topk的token，此外，在选取topk之前还要对attention矩阵进行一维的pooling操作，文中解释这样可以帮助LLM保留更完整的信息，因为只使用窗口选择的内容，会造成生成的信息不全。对token附近token做pooling操作可以补齐上下文，类似高斯模糊。对应下图标号3。
最后，将选出token的 KV Cache 和最后一个窗口中所有Token 的 KV Cache 一起作为 Prompt 的 KV Cache。对应下图标号3。

SnapKV只在 Prefill 阶段筛选出关键的 KV Cache，然后在 Decoding 阶段使用这些稀疏化的 KV Cache，以加速长上下文场景的生成速度。此外，Decoding 阶段也不会再更新 Prompt 的 KV Cache。

具体算法如下。

1.4 针对层特点的稀疏化

LLM的分层架构导致跨层的不同信息提取模式，每层的KV缓存对模型性能的贡献不同。
LLM模型的注意力层有自己的特点，这也影响了KV Cache稀疏化策略。比如，为了估计需要保留多少KV Cache中的键/值，InfiGen作者对每个查询token的注意力权重按降序排序，并累计键token的权重，直到累计权重达到0.9。下图展示了在不同层（Layer 0 和 Layer 18）中，为了达到总注意力权重的0.9所需的键token数量分布情况。直方图的横轴表示键token数量，纵轴表示查询token的数量。

第0层的注意力模式更加分散，不同查询token所需的键token数量差异较大，需要参考更多的键token才能捕获有效的注意力。
第18层绝大部分查询token仅需16至几十个键token即可达到累计权重0.9，只有极少数token需要更多的键token，说明高层的注意力模式更集中，依赖较少的键token即可获得有效注意力。

从这些观察中我们可以知道，每个层对于总注意力其实都有着不同程度的贡献，所以不能直接驱逐在外。而且，不同层的注意力模式差异显著。因此，跨层的统一KV缓存压缩可能不是最优的。
在KV Cache管理中，即需要针对LLM层间信息提取异质性，依据注意力权重的贡献进行区分处理或者累积处理，也需要动态调整每层参与注意力计算的键token数量。通过根据每个组件对预测准确性的重要性智能分配内存资源，从而优化内存利用率，同时最大限度地减少准确性下降。
注：论文“A Survey on Large Language Model Acceleration based on KV Cache Management”是将基于层特点的稀疏化和基于头特点的稀疏化都归结到“KV Cache Budget Allocation”类别中，KV缓存预算分配通过根据每个组件对预测准确性的重要性来智能分配内存资源来解决这一挑战。比如将更多的KV缓存分配给信息更加分散的较低层（前几层），同时减少较高层信息比较集中的KV缓存等。论文给出了具体总结如下图所示。

我们具体看看一些案例。
1.4.1 PyramidKV

论文“PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling”通过研究不同layer间的注意力机制，探索了跨层共享kvcache的效果，发现较低层在输入序列中表现出均匀的注意力分布，而上层则表现出对特定token的集中注意力。因此，PyramidKV采用了金字塔形的内存分配策略，将更多的KV缓存分配给信息更加分散的较低层（前几层），每个KV包含的信息较少，同时减少较高层的KV缓存，使得高层中信息变得集中在较少的关键token中，同时在每一层选择具有高关注值的token。此外，在解码阶段，PyramidInfer通过由注意力值驱动的频繁更新动态维护一组重要token。
洞察

作者研究了Llama模型进行多文档问答的逐层注意力图，发现了注意力层中存在金字塔形信息汇聚模式（Pyramidal Information Funneling）：

在模型的低层（例如第0层）中，注意力得分呈现近似均匀分布，信息比较分散。这表明模型在较低层时从所有可用内容中全局聚合信息，而不会优先关注特定的段落。
当编码信息进行到中间层（6-18）时，逐渐转变为聚焦在段落内部的注意力模式（Localized Attention）。在这个阶段，注意力主要集中在同一文档内的Token上，表明模型在单个段落内进行了段落内部的信息聚合。即模型已经将需要的信息集中到了部分token中。
这种趋势在上层（24-30）继续并加强，而且可以观察到“Attention Sink”和“Massive Activation”现象。

Massive Activation出自论文“Massive activations in large language models”，即在模型的隐藏状态中存在极少数激活值（activations）远大于其他激活值的情况，这些被激活称为“massive activations”。
方案

作者选择通过基于注意力模式动态分配缓存预算来提高压缩效率。针对层的不同特性使用不同的KV Cache机制。

在前几层使用全部的kvcache。因为这些层在聚合全局信息，所以信息更加分散。
越深的层使用越少的kvcache。在这些层中，注意力机制极大地集中在少数几个关键Token上，因此只需要保留这些关键Token就能让输出保持一致并且减少显存占用。而其他的token由于值很低，最后被采样的概率也很低，因此可以考虑丢弃这部分token的cache。

KV Cache最终是一个金字塔形状，即论文中提到的PyrimadKV。一旦为每一层确定了KV缓存预算，PyramidKV在每一个Transformer层中选择根据注意力选择要缓存的KV。最后的部分（$\alpha$个）Token的KV缓存，即Instruction Token，则会在所有Transformer层中保留。

为了量化每个token在生成过程中的重要性，作者测量每个token从指令token中获得的关注程度，并使用此测量值为KV缓存选择重要token。具体实现流程如下：

因为最前面的token和最近邻的token都包含重要信息，对生成效果影响显著，不能舍弃，所以用超参数控制这些token的选择，比如始终保留初始64个token（sink token）和最近的128个token（instruction token）。
确定第一层和最后一层要保留的token总数，比如第一层是全部，最后一层是1/10（超参数控制）的token。那么中间layer需要保存的token数目从第一层到最后一层线性递减。具体计算token是用 instruction token和该token计算注意力分数来确定的。
确定了token数目k之后，就可以进行topk的选择。由于qk矩阵每一行的权重和value值加权平均后得到最终的hidden_status，因此筛选规则就是qk矩阵每一行的权重值。

每一层layer在decoding时，只加载对应数目的kvcache进行计算。

1.4.2 PyramidInfer

PyramidInfer也采用金字塔形的预算分配策略，同时在每一层选择具有高注意力值的token。此外，在解码阶段，PyramidInfer依据注意力值来更新动态维护重要token。
动机

之前很多方案忽视了层与层之间的相互依赖以及预计算过程中巨大的内存消耗。PyramidInfer 的作者发现影响未来生成的关键 KV 的数量逐层减少，并且可以通过注意力权重的一致性来提取这些对将来有关键作用的 KV。基于这些发现，作者提出了 PyramidInfer，通过逐层保留关键上下文来压缩 KV Cache。
下图为 PyramidInfer 与 StreamingLLM 和 H2O 的区别，PyramidInfer 中 KV Cache 会逐层递减，越往后越稀疏。

方案

PyramidInfer 的执行过程如下图所示：

在 Prefill 阶段，PyramidInfer 只保留每层的关键上下文（Pivotal Context, PvC）来压缩 KV Cache。
在 Decoding 阶段，PyramidInfer 根据新的最近的 Token 来更新 PvC。

1.4.3 ZigZagKV

ZigZagKV 其实是在SnapKV上继续做优化，也就是说，它也是输入prompt的KV Cache压缩算法，但是ZigZagKV依据层特点进行了稀疏化，因此放在此处。
动机

为了更好的改进，ZigZagKV 定义如下几个指标：

MBA（Minimum Budget size required to maintain 90% of the totalAttention score），其实就是维持90%的attention score需要的最小KV数。
LMBA（Layer Minimum Budget size to maintain Attention score），由于每个attention head都可以计算出一个MBA值，我们把一层里面所有的attention head的MBA的均值叫做LMBA。LMBA越高也就意味着需要更多的token来维持精度，也就意味着需要分配更多的显存来存储KV Cache。
LMBO（Layer-wise Minimum Budget for Output）：此参数的作用是以输出的相似性这个维度分析一下需要保留多少的KV才能让输出类似。

具体参见下图，其中，$a_i$是A的第i个元素。n是当前KV Cache里面的token数量。

通过实验，作者发现：

不同的模型的对应层的LMBA差距比较大；同一个模型不同层的LMBA差距也很大。
同一个模型不同层在保持输出基本一致的限制下需要的KV数量差距很大。

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

为每一层分配多大的空间存储KV Cache是合理的呢？
一个直观的想法就是每一层按照LMBA计算得到的比例来分配，大层多份，小的少分。具体参见下图标号1，其中B是单层的Budget。这么分其实有一个问题，就是如果某一层的LMBA特别大的时候可能导致某些层分配到的空间特别小，也就意味着这些层基本无法保存任何KV。一个层基本没有KV Cache，效果自然保证不了。所以，我们给每个层保留一个最小大小，分配公式改为下图标号2。这就是ZigZagKV的主要改进点。SanpKV给每一层均匀分配空间。而ZigZagKV自适应地分配空间。
解决了每一层KV Cache的空间分配问题，还有一个问题就是要保存那些K和V。ZigZagKV里面token的选择跟SnapKV非常类似。但是，SnapKV给所有的layer的head是平均分配的，ZigZagKV没有给每个attention head分配同样多的显存，而是根据每个head可以保留绝大部分信息需要的token数来进行分配。

1.4.4 Layer-Condensed KV Cache

LCKV建议只计算和缓存一小部分层的键和值，甚至只计算顶层，然后让底层的查询与保存的键和值配对进行推理。这种方法不仅大大提高了推理速度，降低了内存消耗，而且与现有的内存节省技术正交，能够直接集成以进一步优化。虽然这种机制使下一个注意力计算依赖于前一个注意力的顶层密钥和值，这与Transformer的并行训练相矛盾，但LCKV引入了一种近似训练方法来支持并行训练。
具体而言，Transformer 对模型输入进行层层编码的过程，其实也就是对token信息一步步加工的过程，因此Layer-Condensed KV Cache作者认为i最后一层的信息是最重要最全面的，可以忽略前面几层处理得到的信息，只保留最后一层的KV cache。当生成后续 Token 时，其对应的 KV Cache 都从最后一层取。这样模型不再需要计算以及保存前面这些层token的key和value，也不需要前面几层的$W^K W^V$ 矩阵，极大的节省了计算和内存开销。如下图所示。

1.5 其它方案

1.5.1 针对头特点的稀疏化

针对头特点的稀疏化进行更细粒度的逐头预算分配，它能够在每一层内的单个注意力头之间进行精确的内存分配，提供更灵活和有针对性的优化机会。

AdaKV：基于注意力模式差异为各头分配缓存，通过优化L1损失界限最大化保留信息。
RazorAttention：描述了两类不同的检索头。“回声（echo）头”专注于先前出现的相同token，“归纳（induction）头”关注与当前token重复的之前token。该框架实现了差异化缓存策略，为检索头维护完整的缓存条目，同时将远程注意力压缩为非检索头的合并补偿注意力。
DuoAttention：DuoAttention引入了一种参数化方法来区分两类注意力机制：检索头，对全面的长上下文处理至关重要；流式头，主要处理最近的注意力和注意力吸收器。这种分类是通过学习参数实现的，这些参数可以自动识别需要全注意力的检索头。

1.5.2 投机采样稀疏化

投机采样稀疏化的思路是：通过draft模型生成候选token，然后选取其中重要的token。
论文“TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding”提出了 TriForce，这是一种层次化投机采样系统，主要用于长序列生成。TriForce想解决的是长序列生成背景下的加速问题。在长序列生成设定下，除了模型本身的权重，模型推理的 KV cache 也会占据大量的显存，并在上下文足够长的时候同样成为制约推理速度的关键因素。
因此，TriForce引入了一个只使用部分 KV cache 的 target model 来减小全量的 KV cache 对于推理速度的制约，同时采用多级投机推理策略进一步降低 draft model 的推理开销。TriForce的核心思路如下图所示，其主要分为三个过程，涉及两级投机采样：

Speculation Phase 1（此阶段可以执行多次，直到获得满足条件的多个候选 Token）：
- 使用比较小的 LLaMA-68M 模型作为草稿模型，同时使用 StreamingLLM 的方案快速生成多个候选。
- 使用 LLaMA2-7B-128K 模型作为 Target 模型，采用稀疏 key、value cache 方案进行并行验证。验证后的 Token 作为新的草稿。因为采用的稀疏 key、value cache，因此相应的结果可能是有损的。
Speculation Phase 2（此阶段执行一次）：使用 LLaMA2-7B-128K 模型作为 Target 模型，采用全量 key、value cahe 方案进行并行验证。此阶段可以保证生成的结果无损。

在缓存的使用上，TriForce采用基于检索的方案：

将 key、value 划分为不同的 Chunk。
使用 query 与每个 Chunk 的中心（使用均值作为中心）计算距离。
选择 topk 的 Chunk 作为稀疏 key、value cache 进行 Attention 计算。

1.5.3 聚类稀疏化

为了加速关键注意力的检索，一些研究工作提出了基于索引的方法，以块或集群（cluster ）粒度组织和访问KV缓存，实现了高效的查询和提取操作。InfLLM在块中维护完整的KV缓存，同时通过分层存储策略促进长序列处理。该框架采用CPU-GPU内存编排，在GPU内存中保留基本注意力和当前计算单元，同时将访问频率较低的单元卸载到CPU内存中。为了进一步提高top-k块检索精度，Quest框架提出了一种基于KV缓存块中最小和最大键值的精细块表示方法。
SqueezedAttention在离线阶段采用K-means聚类来对语义相似的密钥进行分组，每组由一个质心表示。在推理过程中，它将输入查询与这些质心进行比较，以识别和加载上下文中语义相关的键。同样，RetrievalAttention使用近似最近邻搜索（ANNS）技术索引KV缓存注意力。此外，EM-LLM动态地将传入的注意力分段为偶发事件。此外，它实现了一种混合检索机制，将语义相似性匹配与时间上下文相结合，以有效地访问相关的KV缓存段。同样，ClusterKV将注意力分组到语义簇中，并在推理过程中选择性地召回它们，从而为LLM实现了高精度和高效率。
我们用ClusterKV为例进行学习。
在如何舍弃K和V的方向的研究主要又细分为两个方向：

Non-recallable KV Cache：直接把一些系统判定为不重要的K和V舍弃掉。如何判定不重要是里面关键点。比如，保留attention sink以及最近的窗口里面的token对应的KV cache，把其它的都扔掉，以此来使得计算和显存都可控。
Recallable KV Cache：实际上，KV Cache里面的某一个token的重要性是动态变化的。一个token也许在前面不重要，但是在后边就变重要了。为了得到更好的结果，我们不能把某个KV直接抛弃掉，而应该根据具体的情况在一些时候抛掉KV，一些时候又把KV重载载入显存继续使用。

下图给出了几种方法的对比，图上的星号表示重要的token。我们可以看出，重要的token其实是动态变化的。

ClusterKV的一个重要观察就是，对于一个q重要的token对应的K的cos距离都比较小。基于这个观察，作者使用K-means聚类的方法把类似的K和对应的V进行聚类(使用K的相似性作为分类依据)，这样每个类都有一个中心向量，计算重要性的时候就使用这个中心向量来计算。如果q和一个中心向量的cos距离(1-cos值)比较小，那么q和这个类里面所有的向量的cos距离都应该比较小，换句话说距离比较小的类都会比较大。所以我们可以把这一个整个类都加载进来做后续计算，并且保证漏掉的重要token尽量少。

聚类方法如下：

先把attention sink部分保留，论文简单地直接把前16个token对应的都保留；
对于prefill阶段的结果进行聚类。作者实验发现，对于32K的context，400个类是最佳权衡；
decoding阶段，每m步做一次聚类。论文中默认情况下每320个token生成完之后新作一次聚类。

在选择类别是，先设置固定大小B，然后按照类别的重要性从最重要的开始取，一直取到所有类里面的元素加起来超过B为止。

0x02 复用

复用也是减少KV序列长度的一个重要手段。在本节，我们主要介绍两种方案：KV Cache合并（Merging）和前缀复用。
2.1 KV Cache合并

面对是否删除不重要的键值对以节省内存，还是保持原样但进行压缩而不删除的问题，尚无定论。删除操作或许能立即缓解内存压力，但可能会潜在地影响模型的性能。相反，更先进的压缩技术则力求在减少内存占用的同时，保持信息的完整性。
KV cache合并（KV cache merging ）探索通过压缩或整合KV缓存，尽可能保留被丢弃token的信息，而不会显著降低模型精度，这可以视为对前面提到的token丢弃策略的扩展。KV缓存合并技术不是采用统一的压缩策略，而是利用层内和层间的固有冗余来动态优化内存利用率。这些方法旨在减小KV缓存的大小，同时保留准确注意力计算所需的关键信息，从而在资源受限的环境中实现高效推理。现有的KV缓存合并策略可分为两种主要方法：层内合并，侧重于在各个层内合并KV缓存，以减少每层的内存使用；跨层合并，针对跨层冗余，以消除不必要的重复。这两种方法都提供了互补的优势，提供了平衡内存节省和模型性能下降的灵活性。KV缓存合并的总结如表4所示。

2.1.1 层间合并

对于层间合并，我们以MiniCache为例进行学习。
MiniCache 会合并中高层相似KV对。具体而言，MiniCache并将将相邻相似层中每个token的KV对分解为向量共享方向向量和幅度，仅存储共享的方向向量、注意力幅度和不可缩放的注意力，以最大限度地提高内存效率。
动机
作者发现，KV Cache 在 LLM 中的深层部分的相邻层之间表现出了高度相似性，可以基于这些相似性对 KV Cache 进行压缩与合并。此外，作者还引入了 Token 保留策略，对高度不同的 KV Cache 不进行合并。并且这种方法可以与其他的 KV Cache 量化方案正交使用。
思路
下图给出了MiniCache和以前方法之间的主要区别。MiniCache的特点是沿着LLM的深度（depth）维度研究KV缓存的层间冗余，这是之前基于层内方法忽略的一个方面。这里，T是指预解码的最后一个时间戳，T+1是指解码的第一个时间戳。

下图是MiniCache的压缩策略和保留策略。

(a)描述了跨层压缩过程。模型从层l和l-1中获取KV缓存，并通过方程式(3)将它们合并到共享状态中。此外，模型也选择不可合并的注意力进行保留，然后将合并的缓存、保留注意力和数量存储在C中。
(b)说明了层l和l-1的恢复过程，其中包括方程式(2)中的数量重新缩放和保留注意力恢复。

算法
下图为MiniCache详细的执行流程。

获取 KV Cache：在 Prefill 阶段，逐层生成 KV Cache。
跨层合并：当到达合并开始层 S 时，将当前层 L 的 KV Cache 与前一层 L-1 的 KV Cache 进行合并，以减少冗余。
缓存：将合并后的 KV Cache 存储起来，以便将来使用。
删除：在 Decoding 阶段，删除不必要的或冗余的 KV Cache，以优化内存使用。
加载和生成：获取所需的 KV Cache，用于生成输出。
恢复：对获取的 KV Cache 应用误差抑制机制，包括 rescaling 和 retention recovery，以最小化合并和压缩过程中引入的误差。
更新：在恢复阶段后，使用最终的 KV Cache 更新共享的 KV Cache。

2.3 层内合并

我们使用DMC来学习层内合并。

在普通 Transformer 中，在每个时间步
来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！

账号		自动登录	找回密码
密码			立即注册