LeanK：面向高效解码的可学习K缓存通道剪枝
LeanK: Learnable K Cache Channel Pruning for Efficient Decoding
AI总结：
LeanK通过可学习的两阶段训练剪枝K缓存通道，实现高达70%的K缓存和16%-18%的V缓存内存减少，并加速解码1.3倍，同时保持准确性。
摘要：
大型语言模型（LLMs）支持长上下文任务，但由于键值（KV）缓存的增长而面临效率挑战。我们提出了LeanK，一种基于学习的方法，通过利用静态通道稀疏性来剪枝不重要的键（K）缓存通道。通过一种新颖的两阶段训练过程，LeanK学习通道级静态掩码，该掩码能够满足特定的稀疏率要求和硬件对齐需求。LeanK在不牺牲准确性的情况下减少了GPU内存并加速了解码。实验表明，K缓存内存减少高达70%，V缓存内存减少16%-18%。自定义解码内核使注意力计算加速1.3倍。我们还通过分析学习到的重要性分布，提供了长上下文推理期间模型通道和注意力头的见解。我们的代码可在https://aka.ms/LeanK获取。
英文摘要：
Large language models (LLMs) enable long-context tasks but face efficiency challenges due to the growing key-value (KV) cache. We propose LeanK, a learning-based method that prunes unimportant key (K) cache channels by leveraging static channel sparsity. With a novel two-stage training process, LeanK learns channel-wise static mask that could satisfy specific sparsity ratio and hardware alignment requirement. LeanK reduces GPU memory and accelerates decoding without sacrificing accuracy. Experiments demonstrate up to 70% K cache and 16%-18% V cache memory reduction. Custom decoding kernel enables 1.3x speedup for attention computation. We also provide insights into model channels and attention heads during long-context inference by analyzing the learned importance distribution. Our code is available at https://aka.ms/LeanK.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
