一、背景
在 LLM 推理中，常常会采用 KV Cache 来缓存之前 Token 的中间结果，以显著减少重复计算，从而降低自回归生成中的延迟。然而，KV Cache 的大小与序列长度成正比，在处理长序列时会面临极大的挑战。尤其当前许多模型开始支持几百 K 甚至几 M 的序列长度，进一步凸显了 KV Cache 的问题，因此很多研究工作致力于降低 KV Cache 的占用。
本文中简单介绍几个最新的工作，包括 SnapKV、YOCO、CLA、Layer-Condensed KV Cache、MiniCache 以及 PyramidInfer，它们都试图降低缓解 KV Cache 的压力。关于 GQA、MQA、DeepSeek MLA 以及量化相关的工作我们已经在之前进行了介绍，这里不再赘述。
二、KV Cache 大小
KV Cache 的大小与模型配置（层数，hidden_size，Attention head 个数等）以及序列长度、Batch Size 成正比。其中单个 Token 对应的 KV Cache 大小与模型配置相关，并且是固定的，这里将其称为单位 KV Cache 计算公式为：
sum_token = (hidden_size / num_attention_heads * num_key_value_heads) * num_hidden_layers * 2(k, v)
而总的 KV Cache 大小为：
sum = sum_token * seq_len * batch_size
batch_size 和 seq_len 越大，KV Cache 越大，如下图所示为 LLaMA2-7B 模型的 batch_size 和 seq_len 对应的 KV Cache 大小（默认 FP16 精度）：
三、SnapKV
[2404.14469] SnapKV: LLM Knows What You are Looking for Before Generation 的核心思路比较简单，如下图 Figure 1 所示，在 Prefill 阶段不是保留所有输入 Token 的 KV Cache，而是采用稀疏化的方式，针对每个 Attention Head 将 Prompt 分为 Prefix 和 Window 两部分；然后，通过 Window 中 Token 与 Prefix 中 Token 的 Attention Score 来选择稀疏化的 Token；最后，将它们的 KV Cache 和 Window 中 Token 的 KV Cache 一起作为 Prompt 的 KV Cache。需要说明的是：每个 Attention Head 中从 Prefix 里挑选的 Token 可能不同。此外，Decoding 阶段也不会再更新 Prompt 的 KV Cache。
SnapKV 在处理 16K Token 的输入时，可以获得 3.6x 的加速，内存效率提升 8.2x。同时在 16 个长序列数据集上保持了与基线模型相当的精度。此外，使用 Huggingface 可以在单个 A100-80GB GPU 上处理 380K 上下文 Token 的任务。
四、YOCO
在 [2405.05254] You Only Cache Once: Decoder-Decoder Architectures for Language Models 中，作者只保留一层全局的 KV Cache。这种设计可以大大降低 GPU 显存的需求，加快 Prefill 阶段。如下图所示，YOCO 模型与常规 Decoder-Only LLM 的区别有几点：
五、CLA
[2405.12981] Reducing Transformer Key-Value Cache Size with Cross-Layer Attention 中作者同样采用 Cross-Attention 机制来降低 KV Cache。不同的是作者并非采用固定层作为 Cross-Attention 的输入，而是采用相邻层，如下图左图所示。最简单的方式就是隔层共享，称作 CLA2，实际也可以每 3 层共享，称作 CLA3，如下图右图所示。此外，这种方法与 MQA 和 GQA 等修改 Attention Head 的方案是兼容的。CLA2 显存减小 2x，CLA3 显存减小 3x。
作者训练 1B 和 3B 参数模型模型实验表明，CLA 相比传统的 MQA 在显存占用、准确性方面可以实现帕累托改进，从而实现更长的序列长度和更大的 Batch Size。（PS：但并不意味着可以优于现在广泛采用的 GQA？）
六、Layer-Condensed KV Cache
在 [2405.10637] Layer-Condensed KV Cache for Efficient Inference of Large Language Models 中，作者同样采用了仅计算和缓存少量层 KV Cache 的方案，从而显著节约显存消耗并提升吞吐量。如下图 Figure 1 所示，仅保留最后一个 Transfomer Block 层的 KV Cache，当生成后续 Token 时其对应的 KV Cache 都从最后一层取。
七、MiniCache
在 [2405.14366] MiniCache: KV Cache Compression in Depth Dimension for Large Language Models 中，作者观察到 KV Cache 在 LLM 中的深层部分的相邻层之间表现出了高度相似性，可以基于这些相似性对 KV Cache 进行压缩。此外，作者还引入了 Token 保留策略，对高度不同的 KV Cache 不进行合并。并且这种方法可以与其他的 KV Cache 量化方案正交使用。
作者在 LLaMA-2、LLaMA-3、Phi-3、Mistral 和 Mixtral 等模型上进行实验，在 ShareGPT 数据集上，采用 4 Bit MiniCache LLaMA–7B 与 FP16 全量 KV Cache 相比实现了 5.02x 的压缩比，推理吞吐提高约 5 倍，显存占用减少 41%，同时性能几乎无损。
如下图 Figure 3 所示为其压缩策略和保留策略：
如下图 Figure A 所示为其详细的执行流程：
在 [2405.12532] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference 中，作者发现影响未来生成的关键 KV 的数量逐层减少，并且可以通过注意力权重的一致性来提取这些关键 KV。基于这些发现，作者提出了 PyramidInfer，通过逐层保留关键上下文来压缩 KV Cache。PyramidInfer 在不牺牲性能的情况下计算更少的 KV，并节约大量显存。实验结果表明，与 Accelerate 相比，PyramidInfer 的吞吐提高了 2.2 倍，KV Cache 的显存占用减少了 54% 以上。
如下图 Figure 2 所示为 PyramidInfer 与 StreamingLLM 和 H2O 的区别，PyramidInfer 中 KV Cache 会逐层递减，越往后越稀疏（PS：如果是这样，那么 Layer-Condensed KV Cache 中只保留最后一层的方案是不是不太合理）：
PyramidInfer 的执行过程如下图 Figure 6 所示：
如下图 Table 1 所示，PyramidInfer 在使用更少 KV Cache 的情况下获得更快的推理速度：
如下图 Figure 11 所示，作者进一步测试了 PyramidInfer 在更多 Batch Size 下的表现，其在 比较小 Batch Size 时几乎没有加速 ，主要是因为减少 KV Cache 还需要一些 ；而在比较大的 Batch Size 能获得更大的加速比。而 Full Cache 当 Batch Size 大于 32 吞吐反而降低：（PS：这个降低不太符合预期，通常来说随着 Batch Size 的增加，计算密度会更高，相应的吞吐也应该更高，而且在 32 左右还远没有到 Compute Bound）。
本文转载自，作者：