SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget

在优化大型语言模型（LLM）的键-值（KV）缓存以节省推理成本方面，被认为是至关重要的。现有的大多数KV缓存压缩算法试图通过利用令牌的不同重要性使令牌序列稀疏化。在这项工作中，我们发现通过识别注意力层的重要性，我们可以从两个维度共同优化KV缓存。根据我们在推理中观察到的层次重要性，我们提出了SqueezeAttention来精确地在运行时优化KV缓存预算的分配，并将三种代表性的令牌稀疏化算法纳入每个层的预算来压缩KV缓存。通过从序列和层的维度优化KV缓存，SqueezeAttention在广泛的LLM和基准测试中实现了约30%到70%的内存减少和最多2.2倍的吞吐量提高。代码可在https://github.com/hetailang/SqueezeAttention上找到。

优化大型语言模型的键值（KV）缓存被认为是节省推理成本的关键。本论文试图通过识别注意力层的重要性，从两个维度共同优化KV缓存，即序列和层次维度。

通过识别注意力层的重要性，提出SqueezeAttention算法，实现在运行时精确优化KV缓存的预算分配，并针对每个层使用三种代表性的令牌稀疏化算法来压缩KV缓存。

SqueezeAttention算法在多种LLM和基准测试中实现了约30%至70%的内存减少和高达2.2倍的吞吐量提高。代码已开源。

相关研究包括以往针对KV缓存的稀疏化算法，以及其他优化LLM推理成本的研究，如模型量化和剪枝等。