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

一、背景

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

本文转载自​​,作者:

您可能还会对下面的文章感兴趣: