CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration
作者: Haoyun Jiang, Haolin Li, Jianwei Zhang, Fei Huang, Qiang Hu, Minmin Sun, Shuai Xiao, Yong Li, Junyang Lin, Jiangchao Yao
分类: cs.LG, cs.AI
发布日期: 2026-08-31
备注: Published at ICML 2025
期刊: Proceedings of the 42nd International Conference on Machine Learning (ICML), PMLR 267:27569-27585, 2025
💡 一句话要点
提出CateKV以解决长上下文LLM推理加速问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 长上下文 推理加速 大型语言模型 KV缓存 注意力机制
📋 核心要点
- 现有大型语言模型在处理长上下文时面临内存需求高和推理延迟大的挑战。
- CateKV通过识别注意力头的顺序一致性，提出了一种混合KV缓存方法，优化了内存和计算效率。
- 实验结果表明，CateKV在保持准确性的同时，内存使用减少最多2.72倍，解码速度提升2.18倍，批处理吞吐量提升3.96倍。
📝 摘要（中文）
大型语言模型（LLMs）在处理长上下文任务方面表现出色，但由于巨大的内存需求和推理延迟，处理长上下文仍然具有挑战性。本文发现某些注意力头在其注意力模式中表现出顺序一致性，这可以通过基于变异系数的算法持续识别。基于这一观察，提出了CateKV，一种混合KV缓存方法，仅保留一致头的关键令牌信息，从而减少KV缓存大小和计算开销，同时在自适应头中保留大部分KV对以确保高准确性。全面评估表明，CateKV在保持与全注意力相当的准确性的同时，内存使用减少了最多2.72倍，单样本输入的解码加速了2.18倍，批处理场景的吞吐量提升了3.96倍。
🔬 方法详解
问题定义：本文旨在解决大型语言模型在处理长上下文时的内存和推理延迟问题。现有方法在长上下文任务中往往需要大量内存，导致推理效率低下。
核心思路：CateKV的核心思路是利用注意力头的顺序一致性，通过保留关键令牌信息来优化KV缓存，从而减少内存占用和计算开销。
技术框架：CateKV方法的整体架构包括两个主要模块：一是基于变异系数的算法用于识别一致头，二是混合KV缓存策略用于存储关键令牌信息和自适应头的KV对。
关键创新：CateKV的主要创新在于识别注意力头的顺序一致性，并基于此设计混合KV缓存方法，与传统方法相比，显著降低了内存使用和计算复杂度。
关键设计：在设计中，采用了变异系数算法来评估注意力头的一致性，并通过动态调整KV缓存的大小和内容，确保在保持高准确性的同时实现推理加速。具体参数设置和损失函数设计在实验部分进行了详细说明。
🖼️ 关键图片
📊 实验亮点
实验结果显示，CateKV在保持与全注意力相当的准确性的同时，内存使用减少最多2.72倍，单样本输入的解码速度提升2.18倍，批处理场景的吞吐量提升高达3.96倍，展现了显著的性能优势。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理中的长文本生成、对话系统以及其他需要处理长上下文的任务。通过优化推理效率，CateKV能够在实际应用中显著提升用户体验，降低计算资源消耗，具有重要的实际价值和未来影响。
📄 摘要（原文）
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.