ClusterKV: Manipulating LLM KV Cache in Semantic Space for Recallable Compression
作者: Guangda Liu, Chengwei Li, Jieru Zhao, Chenqi Zhang, Minyi Guo
分类: cs.LG, cs.AI, cs.PF
发布日期: 2024-12-04 (更新: 2025-06-14)
🔗 代码/项目: GITHUB
💡 一句话要点
ClusterKV：通过语义空间操纵LLM KV缓存，实现可召回的压缩
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 KV缓存压缩 语义聚类 长上下文推理 可召回压缩
📋 核心要点
- 现有KV缓存压缩方法要么永久删除token，要么以文本位置划分的页面粒度召回，导致模型精度和输出质量下降。
- ClusterKV的核心思想是在语义空间中对token进行聚类，并以语义集群的粒度进行召回，从而实现高效且准确的KV缓存压缩。
- 实验结果表明，ClusterKV在32k上下文长度下，仅使用1k-2k KV缓存预算，即可实现可忽略的精度损失，并显著提升推理速度和吞吐量。
📝 摘要（中文）
大型语言模型（LLMs）已广泛应用于各种场景，并且上下文长度正在迅速增加，以处理诸如长文档问答和复杂逻辑推理等任务。然而，长上下文对推理效率提出了重大挑战，包括键值（KV）缓存的高内存成本以及由于大量内存访问而导致的延迟增加。最近的研究提出了压缩KV缓存以近似计算，但这些方法要么永久驱逐token，永远不会为了后续推理而召回它们，要么以文本位置划分的页面粒度召回先前的token。这两种方法都会降低模型准确性和输出质量。为了实现高效且准确的可召回KV缓存压缩，我们引入了ClusterKV，它以语义集群的粒度召回token。我们设计并实现了用于聚类、选择、索引和缓存的高效算法和系统。实验结果表明，ClusterKV在各种具有32k上下文长度的任务中实现了可忽略不计的精度损失，仅使用1k到2k的KV缓存预算，并在延迟方面实现了高达2倍的加速，在解码吞吐量方面实现了2.5倍的提升。与SoTA可召回KV压缩方法相比，ClusterKV展示了更高的模型准确性和输出质量，同时保持或超过了推理效率。我们的代码可在https://github.com/sjtu-zhao-lab/ClusterKV 获得。
🔬 方法详解
问题定义：论文旨在解决长上下文LLM推理中KV缓存占用内存过大，导致推理效率降低的问题。现有KV缓存压缩方法，如直接丢弃token或基于文本位置划分页面进行召回，会严重影响模型精度和生成质量。因此，需要一种既能有效压缩KV缓存，又能保证模型性能的可召回压缩方法。
核心思路：ClusterKV的核心思路是利用token的语义信息进行聚类，将语义相似的token归为一类。在需要压缩KV缓存时，可以根据一定的策略选择性地保留或丢弃某些语义簇。当需要召回被压缩的token时，可以以语义簇为单位进行召回，从而在压缩率和模型性能之间取得平衡。这种基于语义的压缩和召回方式，能够更好地保留上下文信息，减少信息损失。
技术框架：ClusterKV的整体框架包含以下几个主要模块：1) 聚类模块：负责将KV缓存中的token进行语义聚类。2) 选择模块：根据一定的策略，选择需要保留或丢弃的语义簇。3) 索引模块：为每个语义簇建立索引，方便后续的召回操作。4) 缓存模块：负责存储和管理压缩后的KV缓存。在推理过程中，首先查询缓存中是否存在所需的token，如果不存在，则从索引中查找对应的语义簇，并进行召回。
关键创新：ClusterKV的关键创新在于提出了基于语义聚类的可召回KV缓存压缩方法。与传统的基于文本位置的压缩方法相比，ClusterKV能够更好地保留上下文信息，减少信息损失。此外，ClusterKV还设计了高效的聚类、选择、索引和缓存算法，保证了压缩和召回的效率。
关键设计：在聚类模块中，可以使用各种聚类算法，如K-means、DBSCAN等。选择模块可以根据token的重要性、语义簇的大小等因素进行选择。索引模块可以使用哈希表、B树等数据结构。缓存模块可以使用LRU、LFU等缓存策略。具体的参数设置需要根据具体的应用场景进行调整。论文中可能还涉及一些特定的损失函数或网络结构，用于优化聚类效果或提高召回精度，但具体细节未知。
🖼️ 关键图片
📊 实验亮点
实验结果表明，ClusterKV在32k上下文长度下，仅使用1k到2k的KV缓存预算，即可实现与原始模型几乎相同的精度。与现有技术相比，ClusterKV在延迟方面实现了高达2倍的加速，在解码吞吐量方面实现了2.5倍的提升。此外，ClusterKV在模型准确性和输出质量方面也优于其他可召回KV压缩方法。这些结果表明，ClusterKV是一种高效且有效的KV缓存压缩方法。
🎯 应用场景
ClusterKV技术可广泛应用于需要处理长上下文的LLM应用中，例如长文档问答、代码生成、复杂逻辑推理等。通过降低KV缓存的内存占用，可以显著提高LLM的推理效率，降低部署成本，并支持更大规模的上下文窗口。该技术还有助于在资源受限的设备上部署LLM，例如移动设备和嵌入式系统。未来，该技术有望进一步发展，实现更高效、更智能的KV缓存管理。
📄 摘要（原文）
Large Language Models (LLMs) have been widely deployed in a variety of applications, and the context length is rapidly increasing to handle tasks such as long-document QA and complex logical reasoning. However, long context poses significant challenges for inference efficiency, including high memory costs of key-value (KV) cache and increased latency due to extensive memory accesses. Recent works have proposed compressing KV cache to approximate computation, but these methods either evict tokens permanently, never recalling them for later inference, or recall previous tokens at the granularity of pages divided by textual positions. Both approaches degrade the model accuracy and output quality. To achieve efficient and accurate recallable KV cache compression, we introduce ClusterKV, which recalls tokens at the granularity of semantic clusters. We design and implement efficient algorithms and systems for clustering, selection, indexing and caching. Experiment results show that ClusterKV attains negligible accuracy loss across various tasks with 32k context lengths, using only a 1k to 2k KV cache budget, and achieves up to a 2$\times$ speedup in latency and a 2.5$\times$ improvement in decoding throughput. Compared to SoTA recallable KV compression methods, ClusterKV demonstrates higher model accuracy and output quality, while maintaining or exceeding inference efficiency. Our code is available at https://github.com/sjtu-zhao-lab/ClusterKV.