arXiv:2505.22913cs.LG2025-05NeurIPS被引 10

通过无结构稀疏性压缩注意力缓存,实现70%稀疏度下不降精度

Mustafar: Promoting Unstructured Sparsity for KV Cache Pruning in LLM Inference

  • 采用逐令牌幅度剪枝,对键值缓存实现无结构稀疏
  • 支持最高70%稀疏度,推理吞吐提升2.23倍,上下文长度更长
  • 自研稀疏注意力核+位图格式,无需微调即可加速解码

我们证明了无结构稀疏性可显著提升大语言模型推理中键值(KV)缓存的压缩效果,实现高达70%的稀疏度而无需微调且不损失精度。系统性探索表明,基于每令牌幅度的剪枝在键和值缓存上均有效,优于以往的结构化剪枝方案。键缓存受益于显著的异常值元素,而值缓存虽分布均匀,仍可通过简单幅度剪枝获益。针对解码阶段高内存开销的瓶颈,我们采用位图格式存储稀疏缓存,并设计专用注意力内核,可直接在压缩缓存上进行计算,有效加速内存密集型操作。该方法将KV缓存压缩至密集推理的45%,使上下文长度延长,吞吐量最高提升2.23倍。相关代码已开源。

原文摘要 · Abstract (English)

We demonstrate that unstructured sparsity significantly improves KV cache compression for LLMs, enabling sparsity levels up to 70% without compromising accuracy or requiring fine-tuning. We conduct a systematic exploration of pruning strategies and find per-token magnitude-based pruning as highly effective for both Key and Value caches under unstructured sparsity, surpassing prior structured pruning schemes. The Key cache benefits from prominent outlier elements, while the Value cache surprisingly benefits from a simple magnitude-based pruning despite its uniform distribution. KV cache size is the major bottleneck in decode performance due to high memory overhead for large context lengths. To address this, we use a bitmap-based sparse format and a custom attention kernel capable of compressing and directly computing over compressed caches pruned to arbitrary sparsity patterns, significantly accelerating memory-bound operations in decode computations and thereby compensating for the overhead of runtime pruning and compression. Our custom attention kernel coupled with the bitmap-based format delivers substantial compression of KV cache upto 45% of dense inference and thereby enables longer context length and increased tokens/sec throughput of upto 2.23x compared to dense inference. Our pruning mechanism and sparse attention kernel is available at https://github.com/dhjoo98/mustafar.

稀疏化注意力缓存推理优化大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。