大模型压缩KV缓存新突破，中科大提出自适应预算分配，工业界已落地vLLM框架

这篇文章讨论了中国科学技术大学（USTC）的研究团队在大型模型压缩领域取得的重大进展，特别是在KV缓存压缩方面的突破。该团队提出了一种名为Ada-KV的创新算法，利用自适应预算分配来优化KV缓存的淘汰过程，从而提高推理效率。这种方法摆脱了传统方法在所有注意力头上均匀分配压缩预算的做法，而是根据不同头部的注意力模式分配压缩预算。这项研究不仅在学术界引起了讨论，还在工业界得到了应用，Cloudflare Workers AI团队将改进措施整合到了vLLM框架中，这是一个行业标准的部署工具。文章还提到，Ada-KV的代码已在GitHub上完全开源，鼓励在该领域进行进一步研究和应用。对于那些对机器学习、模型优化和AI推理任务中高效资源利用感兴趣的人来说，这一发展尤为值得关注。- 中科大提出自适应预算分配算法来优化KV缓存的驱逐过程，以提高推理效率。 - 大模型在自回归生成过程中，每生成一个新token都需要将对应的KV矩阵存储下来，导致缓存随着生成序列长度的增加而急剧膨胀，引发内存和I/O延迟问题。 - 中科大团队提出的自适应预算分配能够根据不同注意力头的特性差异，提高压缩质量。 - Cloudflare workers AI团队将该方法落地于工业部署常用的vLLM框架中，并开源全部代码。 - 实验结果显示，自适应预算分配的压缩方法优于原有的均匀预算分配压缩方法。 - 该方法结合了展平的KV Cache管理布局和定制的CUDA kernel，实现了高效的Cache更新管理。 - Cloudflare公司的Workers AI团队基于Paged Attention重新实现了该算法，并将其落地于实际部署使用的推理框架vLLM中。