通过学习剪枝不重要键缓存通道，显著降低大模型解码内存占用。
LeanK: Learnable K Cache Channel Pruning for Efficient Decoding
- 基于静态通道稀疏性，学习通道级掩码实现可配置压缩率。
- 最高减少70%键缓存内存，16%-18%值缓存内存，加速解码1.3倍。
- 适合长文本生成、资源受限部署场景，支持硬件对齐优化。
大型语言模型（LLMs）虽能处理长上下文任务，但随上下文增长导致键值（KV）缓存效率下降。本文提出LeanK，一种基于学习的键（K）缓存通道剪枝方法，利用静态通道稀疏性。通过新颖的两阶段训练过程，LeanK学习通道级静态掩码，满足特定稀疏率与硬件对齐需求。实验表明，该方法可实现最高70%的键缓存和16%-18%的值缓存内存减少，自定义解码内核使注意力计算提速1.3倍。同时，通过分析学习到的重要度分布，揭示了长上下文推理中模型通道与注意力头的行为特征。代码已开源于https://aka.ms/LeanK。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。