SqueezeAttention: 通过逐层最优预算在LLM推理中对KV缓存进行二维管理
机器学习
2024-10-11 v2 计算与语言
摘要
优化大型语言模型（LLM）的键值（KV）缓存被认为对节省推理成本至关重要。现有的KV缓存压缩算法大多试图利用令牌的不同重要性来稀疏化令牌序列。然而，这些方法大多平等对待所有层，为每一层分配相同的KV预算。这种方法是次优的，因为某些层可能对输入令牌不太敏感，但仍然获得与其他层相同的预算。在这项工作中，我们发现通过识别注意力层的重要性，我们可以从两个维度（即序列维度和层维度）联合优化KV缓存。基于我们对推理中层重要性的观察，我们提出了SqueezeAttention，以在运行时精确优化各层之间的KV缓存预算分配，然后结合三种代表性的序列级算法，为每一层使用其自己的预算压缩KV缓存。具体来说，我们首先通过计算自注意力层前后输入提示差异的余弦相似度来衡量每层的重要性。基于这种相似性，我们将层分为两组，并相应调整它们的KV预算。通过从序列和层两个维度优化KV缓存，SqueezeAttention在广泛的LLM和基准测试中实现了约30%到70%的内存减少和高达2.2倍的吞吐量提升。代码可在https://github.com/hetailang/SqueezeAttention获取。
引用
@article{arxiv.2404.04793,
title = {SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget},
author = {Zihao Wang and Bin Cui and Shaoduo Gan},
journal= {arXiv preprint arXiv:2404.04793},
year = {2024}
}