发现注意力头具有一致性,用新方法压缩缓存,提速降存。
CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration

- 基于变异系数识别可复用的注意力模式,仅保留关键记忆。
- 单样本推理加速2.18倍,内存减少2.72倍,批量吞吐提升3.96倍。
- 适合长文本生成、低资源部署等场景,兼顾速度与精度。
大型语言模型在处理长上下文任务时表现出强大能力,但因内存需求高和推理延迟大而面临挑战。本文发现部分注意力头在注意力模式上具有序列一致性,可通过基于变异系数的算法持续识别。受此启发,提出CateKV——一种混合键值缓存方法,对一致头仅保留关键标记信息,从而减少缓存大小与计算开销;对自适应头则保留大部分键值对以保障精度。实验表明,在保持与全注意力相当准确率的前提下,单样本输入下内存使用减少2.72倍,解码速度提升2.18倍;批量场景中吞吐量提高3.96倍。该方法可与现有加速技术兼容并扩展。
原文摘要 · Abstract (English)
Large language models (LLMs) have demonstrated strong capabilities in handling long-context tasks, but processing such long contexts remains challenging due to the substantial memory requirements and inference latency. In this work, we discover that certain attention heads exhibit sequential consistency in their attention patterns, which can be persistently identified using a coefficient-of-variation-based algorithm. Inspired by this observation, we propose CateKV, a hybrid KV cache method that retains only critical token information for consistent heads, thereby reducing KV cache size and computational overhead, while preserving the majority of KV pairs in adaptive heads to ensure high accuracy. We show the unique characteristics of our algorithm and its extension with existing acceleration methods. Comprehensive evaluations on long-context benchmarks show that, while maintaining accuracy comparable to full attention, CateKV reduces memory usage by up to $2.72\times$ and accelerates decoding by $2.18\times$ in single-sample inputs, and boosts throughput by $3.96\times$ in batch scenarios.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。