arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2405.12532cs.CL

PyramidInfer:用于高吞吐量LLM推理的金字塔KV缓存压缩

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

Shanghai Jiao Tong University(上海交通大学);Xiaohongshu Inc.(小红书科技有限公司);South China University of Technology(华南理工大学)

Dongjie Yang, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, Hai Zhao

AI总结:

针对LLM推理中KV缓存占用GPU内存大、层间依赖被忽略的问题,提出PyramidInfer方法,通过逐层保留关键上下文压缩KV缓存,在保持性能的同时将吞吐量提升2.2倍、GPU内存减少超54%。

摘要:

大型语言模型(LLMs)展现了卓越的理解能力,但在推理过程中面临GPU内存使用的挑战,阻碍了其在聊天机器人等实时应用中的可扩展性。为加速推理,我们将计算得到的键和值(KV cache)存储在GPU内存中。现有方法通过修剪预计算的KV缓存来减少内存占用。然而,这些方法忽略了层间的依赖关系以及预计算中的巨大内存消耗。为探究这些不足,我们发现影响未来生成的关键键值对数量逐层递减,且可通过注意力权重的一致性提取这些关键信息。基于此发现,我们提出PyramidInfer,一种通过逐层保留关键上下文来压缩KV缓存的方法。PyramidInfer通过计算更少的键值对节省大量内存,且不牺牲性能。实验结果表明,与Accelerate相比,PyramidInfer将吞吐量提升2.2倍,同时KV缓存的GPU内存减少超过54%。

英文摘要:

Large Language Models (LLMs) have shown remarkable comprehension abilities but face challenges in GPU memory usage during inference, hindering their scalability for real-time applications like chatbots. To accelerate inference, we store computed keys and values (KV cache) in the GPU memory. Existing methods study the KV cache compression to reduce memory by pruning the pre-computed KV cache. However, they neglect the inter-layer dependency between layers and huge memory consumption in pre-computation. To explore these deficiencies, we find that the number of crucial keys and values that influence future generations decreases layer by layer and we can extract them by the consistency in attention weights. Based on the findings, we propose PyramidInfer, a method that compresses the KV cache by layer-wise retaining crucial context. PyramidInfer saves significant memory by computing fewer keys and values without sacrificing performance. Experimental results show PyramidInfer improves 2.2x throughput compared to Accelerate with over 54% GPU memory reduction in KV cache.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑