PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference

大型语言模型（LLMs）展现了出色的理解能力，但在推理过程中面临GPU内存使用方面的挑战，这阻碍了它们在聊天机器人等实时应用中的可扩展性。为了加速推理，我们将计算出的键和值（KV缓存）存储在GPU内存中。现有方法研究了KV缓存压缩以通过修剪预计算的KV缓存来减少内存。然而，它们忽略了层间依赖关系和预计算中的巨大内存消耗。为了探索这些缺陷，我们发现影响未来生成的关键键和值的数量逐层递减，我们可以通过注意力权重的一致性提取它们。基于这些发现，我们提出了PyramidInfer方法，该方法通过逐层保留关键上下文来压缩KV缓存。PyramidInfer通过计算更少的键和值而不损失性能来节省大量内存。实验结果显示，与Accelerate相比，PyramidInfer的吞吐量提高了2.2倍，KV缓存的GPU内存减少了超过54%。

解决GPU内存使用限制对于实时应用如聊天机器人的影响，通过压缩KV缓存来加速推理过程，同时保持性能。

通过逐层提取关键上下文信息，实现分层压缩KV缓存，从而减少预计算的键值对数量，节省GPU内存，并提高推理效率。

PyramidInfer方法在保持性能的同时，相比于现有方法，在KV缓存中节省了54%的GPU内存，并提高了2.2倍的推理吞吐量。实验使用了公开数据集，并且开源了代码。

相关研究包括对KV缓存压缩的研究，如DeepSpeed和ZeRO，以及对于模型推理加速的研究，如TensorRT和ONNX Runtime等。