The deployment of long-context Large Language Models (LLMs) poses significant challenges due to the intense computational cost of self-attention and the substantial memory overhead of the Key-Value Cache (KV Cache). In this paper, we introduce HieraSparse, a hierarchical KV Cache compression framework with acceleration kernels that leverage GPU sparse tensor cores to speed up semi-structured KV Cache attention for both the prefill and decode phases. With the hierarchical design, our method allows for a flexible quality-sparsity trade-off and successfully converts sparsity into efficiency. Compared to the state-of-the-art decode method that utilizes unstructured sparsity, HieraSparse achieves $\mathbf{1.2\times}$ KV compression ratio and $\mathbf{4.57\times}$ attention speedup at the same sparsity level. Furthermore, we extended the semi-structured KV Cache pruning to the prefill stage, which demonstrated up to $\mathbf{1.85\times}$ attention speedup at the highest sparsity. Lastly, we evaluate the generation quality of HieraSparse with a simple magnitude-based pruning method, and the results show that $\mathbf{1.37\times}$ prefill speedup and $\mathbf{1.77\times}$ decode speedup can be achieved without significant quality drop. The codebase can be found at https://github.com/psl-ntu/HieraSparse.


翻译:长上下文大语言模型(LLMs)的部署面临严峻挑战,这源于自注意力机制的高昂计算成本以及键值缓存(KV Cache)显著的内存开销。本文提出HieraSparse——一种层级化KV缓存压缩框架,其配套加速内核利用GPU稀疏张量核心加速预填充和解码阶段的半结构化KV缓存注意力运算。通过层级化设计,我们的方法实现了灵活的质量-稀疏度权衡,并成功将稀疏性转化为效率。与当前采用非结构化稀疏性的最优解码方法相比,HieraSparse在相同稀疏度水平下实现了$\mathbf{1.2\times}$的KV压缩比和$\mathbf{4.57\times}$的注意力加速。此外,我们将半结构化KV缓存剪枝扩展至预填充阶段,在最稀疏场景下展现了高达$\mathbf{1.85\times}$的注意力加速。最终,我们采用基于幅度的简单剪枝方法评估HieraSparse的生成质量,结果表明在无显著质量损失的前提下可实现$\mathbf{1.37\times}$的预填充加速和$\mathbf{1.77\times}$的解码加速。代码库可访问https://github.com/psl-ntu/HieraSparse获取。

0
下载
关闭预览

相关内容

TransMLA:多头潜在注意力(MLA)即为所需
专知会员服务
23+阅读 · 2025年2月13日
可视化循环神经网络的注意力机制
论智
22+阅读 · 2018年9月23日
Hierarchical Imitation - Reinforcement Learning
CreateAMind
19+阅读 · 2018年5月25日
【干货】基于Keras的注意力机制实战
专知
59+阅读 · 2018年5月4日
深度学习中的注意力机制
人工智能头条
16+阅读 · 2017年11月2日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
1+阅读 · 2014年12月31日
国家自然科学基金
11+阅读 · 2012年12月31日
Arxiv
0+阅读 · 3月2日
VIP会员
相关主题
最新内容
锻造未来士兵:外骨骼、基因工程与赛博格
专知会员服务
0+阅读 · 15分钟前
《无人机蜂群通信技术研究》50页
专知会员服务
1+阅读 · 32分钟前
战力倍增器:自主武器系统与乌克兰及加沙冲突
人工智能赋能战场情报:提速决策进程
专知会员服务
3+阅读 · 7月17日
《拥抱新兴技术:面向未来军官的教育革新》
专知会员服务
7+阅读 · 7月17日
相关VIP内容
TransMLA:多头潜在注意力(MLA)即为所需
专知会员服务
23+阅读 · 2025年2月13日
相关基金
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
1+阅读 · 2014年12月31日
国家自然科学基金
11+阅读 · 2012年12月31日
Top
微信扫码咨询专知VIP会员