Palu: 基于低秩投影压缩 KV 缓存
人工智能
2024-11-05 v2 机器学习
摘要
后训练 KV 缓存压缩方法通常要么抽样一部分有效 token，要么将数据量化到更低的数值位宽。然而，这些方法无法利用 KV 张量隐藏维度中的冗余。本文提出了一种称为 Palu 的隐藏维度压缩方法，作为一种利用低秩投影来减少 LLM 推理时内存使用的 KV 缓存压缩框架。Palu 将线性层分解为低秩矩阵，缓存压缩后的中间状态，并在实时重建完整的 key 和 value。为提高准确度、压缩率和效率，Palu 还包括 (1) 中等粒度的低秩分解方案、(2) 高效的秩搜索算法、(3) 低秩感知量化兼容性增强、以及 (4) 优化的 GPU 内核与算子融合。大量实验表明，Palu 使用流行的 LLM 可将 KV 缓存压缩 50%，同时保持强大的准确度，并在基于 RoPE 的注意力模块上实现最高 1.89 倍的性能提升。当与量化结合使用时，Palu 天然的量化友好设计可在保持准确度的同时节省额外内存，超过仅量化方法，实现基于 RoPE 的注意力最高可达 2.91 倍的加速。此外，它在准确度方面与仅量化方法相当，甚至更好（最高可降低 1.19 个 perplexity）。这些结果表明 Palu 在有效应对 KV 缓存在 LLM 推理中所致的效率和内存挑战方面具有卓越的能力。我们的代码已公开提供：https://github.com/shadowpa0327/Palu
引用
@article{arxiv.2407.21118,
title = {Palu: Compressing KV-Cache with Low-Rank Projection},
author = {Chi-Chih Chang and Wei-Cheng Lin and Chien-Yu Lin and Chong-Yan Chen and Yu-Fang Hu and Pei-Shuo Wang and Ning-Chi Huang and Luis Ceze and Mohamed S. Abdelfattah and Kai-Chiang Wu},
journal= {arXiv preprint arXiv:2407.21118},
year = {2024}
}