通过聚类压缩长文本记忆,让大模型推理更快更省内存。
CentroidKV: Efficient Long-Context LLM Inference via KV Cache Clustering
- 将序列分块后按相似性聚类,用中心向量替代冗余记忆。
- 内存占用降低75%,推理速度提升1.92倍,吞吐量最高增4倍。
- 适合部署长文本生成场景,对算力要求低且效果稳定。
具备长上下文窗口的大语言模型在处理复杂任务中日益普及,但其所需的大量键值(KV)缓存带来了显著的部署挑战。现有方法要么丢弃未来生成所需的关键信息,要么因计算开销高而效率提升有限。本文提出CentroidKV,一种高效的在线KV缓存聚类框架。基于关键状态在序列维度上高度相似的观察,我们先将序列分块,提出分块软匹配机制,在每块内采用交替划分策略并根据相似性识别聚类,再将每簇内的KV缓存合并为单一中心向量。此外,我们提供了计算复杂度及块内划分策略最优性的理论分析。在多种模型和长上下文基准上的实验表明,CentroidKV可实现高达75%的KV缓存内存缩减,同时保持相近的模型性能;且仅需极小计算开销,推理解码阶段加速达1.92倍,服务吞吐量最高提升4倍。
原文摘要 · Abstract (English)
Large language models (LLMs) with extended context windows have become increasingly prevalent for tackling complex tasks. However, the substantial Key-Value (KV) cache required for long-context LLMs poses significant deployment challenges. Existing approaches either discard potentially critical information needed for future generations or offer limited efficiency gains due to high computational overhead. In this paper, we introduce CentroidKV, a simple yet effective framework for online KV cache clustering. Our approach is based on the observation that key states exhibit high similarity along the sequence dimension. To enable efficient clustering, we divide the sequence into chunks and propose Chunked Soft Matching, which employs an alternating partition strategy within each chunk and identifies clusters based on similarity. CentroidKV then merges the KV cache within each cluster into a single centroid. Additionally, we provide a theoretical analysis of the computational complexity and the optimality of the intra-chunk partitioning strategy. Extensive experiments across various models and long-context benchmarks demonstrate that CentroidKV achieves up to 75% reduction in KV cache memory usage while maintaining comparable model performance. Moreover, with minimal computational overhead, CentroidKV accelerates the decoding stage of inference by up to $1.92\times$ and increases the serving throughput by up to $4\times$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。