英伟达: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令牌,通过筛选策略,保证正确性的同时实现压缩方法。

🧩附录


相关推荐
EAI-Robotics7 分钟前
计算机视觉的数据收集与标注
人工智能·计算机视觉
IT_陈寒14 分钟前
SpringBoot性能调优实战:5个让接口响应速度提升300%的关键配置
前端·人工智能·后端
东临碣石8221 分钟前
【AI论文】OpenGPT-4o-Image:面向高级图像生成与编辑的综合性数据集
人工智能
无风听海43 分钟前
神经网络之激活函数Softmax
人工智能·深度学习·神经网络
无风听海43 分钟前
神经网络之理解温度对Softmax的影响
人工智能·深度学习
关关长语1 小时前
(三) Dotnet为AI控制台对话添加依赖注入等集成
人工智能·ai·ai dotnet
格林威1 小时前
不同光谱的工业相机有哪些?能做什么?
图像处理·人工智能·深度学习·数码相机·计算机视觉·视觉检测
CiLerLinux2 小时前
第三十八章 ESP32S3 SPIFFS 实验
图像处理·人工智能·单片机·嵌入式硬件
苏苏susuus2 小时前
NLP:迁移学习基础讲解
人工智能·自然语言处理·迁移学习
SEO-狼术2 小时前
Oxygen AI Positron Assistant Enterprise
人工智能·.net