arXiv:2412.03213cs.LGcs.AI2024-12被引 68

通过语义聚类压缩大模型缓存,实现高效可召回的长文本推理。

ClusterKV: Manipulating LLM KV Cache in Semantic Space for Recallable Compression

  • 按语义聚类而非位置分页召回缓存,提升信息保留精度。
  • 32k上下文下仅用1k~2k缓存空间,准确率几乎无损失。
  • 适合需要长文本处理的AI应用,如文档问答与逻辑推理。

大语言模型在各类应用中广泛部署,上下文长度持续增长以应对长文档问答和复杂逻辑推理任务。然而,长上下文带来显著的推理效率挑战,包括键值缓存(KV cache)的高内存开销和因频繁内存访问导致的延迟增加。现有压缩方法要么永久丢弃标记,无法召回;要么按文本位置分页召回,降低模型准确性和输出质量。为实现高效且可召回的KV缓存压缩,本文提出ClusterKV,以语义聚类为粒度进行缓存召回。设计并实现了高效的聚类、选择、索引与缓存算法与系统。实验表明,当上下文长度达32k时,使用1k~2k的缓存预算,准确率损失可忽略不计,并实现高达2倍的延迟降低和2.5倍的解码吞吐提升。相比现有最优可召回压缩方法,ClusterKV在保持或超越推理效率的同时,展现出更高的模型准确率与输出质量。代码已开源:https://github.com/sjtu-zhao-lab/ClusterKV。

原文摘要 · Abstract (English)

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.

大模型推理缓存压缩语义聚类长上下文

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。