📖标题：RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
🌟摘要
🔸基于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令牌，通过筛选策略，保证正确性的同时实现压缩方法。
🧩附录
