⏶11
LeanK: 可学习K缓存通道剪枝实现高效解码
发表
由 
zhangyik21 提交
zhangyik21 提交作者: Yike Zhang, 
Zhiyuan He, Huiqiang Jiang, Chengruidong Zhang, Yuqing Yang, Jianyong Wang, Lili Qiu
Yike Zhang, Zhiyuan He, Huiqiang Jiang, Chengruidong Zhang, Yuqing Yang, Jianyong Wang, Lili Qiu摘要
大型语言模型（LLM）支持长上下文任务，但由于不断增长的键值（KV）缓存而面临效率挑战。我们提出了 LeanK，一种基于学习的方法，通过利用静态通道稀疏性来修剪不重要的键（K）缓存通道。通过新颖的两阶段训练过程，LeanK 学习通道级别的静态掩码，该掩码可以满足特定的稀疏性比率和硬件对齐要求。LeanK 减少 GPU 内存并加速解码，而不会牺牲准确性。实验表明 K 缓存减少了高达 70%，V 缓存减少了 16%-18%。自定义解码内核使注意力计算速度提高了 1.3 倍。我们还通过分析学习到的重要性分布，提供了对长上下文推理过程中模型通道和注意力头的见解。我们的代码可在 https://aka.ms/LeanK 获取。
大型语言模型（LLMs）能够处理长上下文任务，但由于键值（KV）缓存不断增长而面临效率挑战。LeanK提出了一种新颖的KV缓存剪枝方法，该方法利用QK向量沿通道（head_dim）维度的静态稀疏性。通过新颖的两阶段训练过程，获得一个静态的、通道维度的剪枝掩码。实验结果表明，K缓存减少了高达70%，V缓存减少了16-18%，对端到端模型性能影响极小。
该方法为KV缓存压缩引入了一个新方向，与现有方法（如令牌级压缩和量化）正交。此外，LeanK观察到每个注意力头中QK向量的通道范数分布与头的检索能力相关。这一见解为增强LLM的长上下文理解能力开辟了有前景的途径。