英伟达:LLM两阶段KV缓存压缩

📖标题:RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression

🌐来源:arXiv, 2502.14051

🌟摘要

🔸基于Transformer的大型语言模型在解码阶段严重依赖于KV缓存来有效地处理扩展上下文。然而,KV缓存的大小与输入长度成比例增长,随着解码的进行,会给内存带宽和容量带来负担。

🔸为了应对这一挑战,我们提出了RocketKV,这是一种无需训练的KV缓存压缩策略,专门用于在解码阶段减少KV缓存的内存带宽和容量需求。RocketKV包含两个连续的阶段。在第一阶段,它使用SnapKV++对输入序列令牌执行粗粒度KV缓存驱逐,这是一种在SnapKV基础上改进的方法,引入了自适应池大小,并与分组查询注意力完全兼容。在第二阶段,它采用混合注意力方法进行细粒度top-k稀疏注意力,通过利用头部和序列维数减少来近似注意力得分。结合这两个阶段,RocketKV实现了显著的KV缓存提取带宽和存储节省,同时保持了与完整KV缓存注意力相当的准确性。

🔸我们发现,与完整的KV缓存基线相比,RocketKV在NVIDIA H100 GPU上的解码阶段提供了高达3倍的端到端加速,峰值内存减少了31%,同时在各种长上下文任务上实现了可忽略的精度损失。

🛎️文章简介

🔸研究问题:长上下文大语言模型(LLM)推理过程中,键值缓存(KV cache)在解码阶段成为主要瓶颈,特别是在内存带宽和容量需求方面。

🔸主要贡献:论文提出了一种名为RocketKV的两阶段KV缓存压缩方法,通过结合永久KV令牌驱逐和动态KV令牌选择,显著加速了LLM推理的解码阶段,同时在保持模型准确性的情况下降低了内存带宽和容量的需求。

📝重点思路

🔸SnapKV++:第一阶段的粗粒度KV缓存驱逐方法,针对输入提示进行有效的KV令牌保留,通过聚合每个注意力组的得分来选择关键KV令牌。

🔸混合注意力(Hybrid Attention):第二阶段的细粒度动态KV令牌选择方法,通过同时利用序列维度和头维度的稀疏性来改进KV令牌的选择,增强预测准确性。

🔸实验评估:在多个长上下文基准上(如LongBench、RULER等)进行对比实验,评估RocketKV与其他方法在准确性、速度和内存使用方面的性能。

🔎分析总结

🔸RocketKV在低令牌预算(如256或512)下,能够显著降低内存带宽和容量使用,同时保持与全KV缓存接近的准确性。

🔸在不同模型和基准测试中,RocketKV在多个设置下表现出优越的性能,尤其在低令牌预算时,相比其他方法具有更小的准确性损失。

🔸实验结果显示,在NVIDIA H100 GPU上,RocketKV在解码阶段实现了高达3倍的速度提升和31%的峰值内存减少。

💡个人观点

论文的核心在于选择重要的KV令牌,通过筛选策略,保证正确性的同时实现压缩方法。

🧩附录


相关推荐
优软轻创-拓客私域10 小时前
数字权益市场爆发:如何通过权益数卡选对优质货源
大数据·人工智能
EllenLiu11 小时前
从 Transformer 理论到文本分类:BERT 微调实战总结
人工智能
绿算技术11 小时前
绿算GP Spark引爆关注,成为AI工厂存储利器
大数据·人工智能·spark
树欲静而风不止8611 小时前
全星QMS:破解高端制造质量难题的一体化数字平台
人工智能
新智元11 小时前
马斯克「世界模拟器」首曝,1 天蒸馏人类 500 年驾驶经验!擎天柱同脑进化
人工智能·openai
新智元11 小时前
LeCun 怒揭机器人最大骗局,坦白 Llama 与我无瓜!
人工智能·openai
俞凡11 小时前
10 分钟搞定神经网络
人工智能
北极的树11 小时前
Claude Agent SDK实战:打造开源版DeepWiki
人工智能
JJJJ_iii11 小时前
【机器学习08】模型评估与选择、偏差与方差、学习曲线
人工智能·笔记·python·深度学习·学习·机器学习
phoenix@Capricornus11 小时前
样本与样本值
人工智能·机器学习·概率论