的最新上班 和 6 MiniCache KV Cache 种优化 等 LLM PyramidInfer

一、背景

在 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）。

本文转载自​​，作者：