用语义分组优化长文本生成的缓存管理,显著降低内存占用
IceCache: Memory-efficient KV-cache Management for Long-Sequence LLMs

- 通过语义聚类将相关词元集中存储,提升缓存效率
- 256词元预算下保持99%原模型准确率,仅需其他方法25%缓存预算
- 适合长文本生成、链式推理等需要长时间序列的任务
键值(KV)缓存通过存储中间注意力状态加速大语言模型的自回归生成,避免重复计算。但其内存开销随序列长度线性增长,在资源受限设备上常造成严重瓶颈。以往工作尝试将部分缓存移至CPU,但依赖粗略的词元选择,导致长文本生成任务中性能下降。本文提出IceCache,结合语义词元聚类与分页注意力机制,利用层次化动态更新的数据结构,将语义相关的词元组织在连续内存区域中,实现更高效的词元选择和跨CPU-GPU传输的带宽利用。在LongBench上的实验表明,使用256词元缓存预算时,IceCache可保持99%的原始准确率;相比其他卸载方法,其延迟与准确率相当甚至更优,且仅需25%的缓存词元预算,充分证明其在长序列场景中的有效性。
原文摘要 · Abstract (English)
Key-Value (KV) cache plays a crucial role in accelerating inference in large language models (LLMs) by storing intermediate attention states and avoiding redundant computation during autoregressive generation. However, its memory footprint scales linearly with sequence length, often leading to severe memory bottlenecks on resource-constrained hardware. Prior work has explored offloading KV cache to the CPU while retaining only a subset on the GPU, but these approaches often rely on imprecise token selection and suffer performance degradation in long-generation tasks such as chain-of-thought reasoning. In this paper, we propose a novel KV cache management strategy, IceCache, which integrates semantic token clustering with PagedAttention. By organizing semantically related tokens into contiguous memory regions managed by a hierarchical, dynamically updatable data structure, our method enables more efficient token selection and better utilization of memory bandwidth during CPU-GPU transfers. Experimental results on LongBench show that, with a 256-token budget, IceCache maintains 99% of the original accuracy achieved by the full KV cache model. Moreover, compared to other offloading-based methods, IceCache attains competitive or even superior latency and accuracy while using only 25% of the KV cache token budget, demonstrating its effectiveness in long-sequence scenarios. The code is available on our project website at https://yuzhenmao.github.io/IceCache/.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。