论文

CompressKV:语义检索引导的KV缓存压缩实现资源高效长上下文LLM推理

CompressKV: Semantic-Retrieval-Guided KV-Cache Compression for Resource-Efficient Long-Context LLM Inference

模型推理KV Cache

摘要

长上下文大语言模型(LLM)推理日益受键值(KV)缓存的内存占用与解码成本约束——限制在资源受限硬件上的可持续部署。既有KV缓存驱逐方法通常在GQA式LLM的全部头上应用启发式token评分。这些方法忽视注意力头的不同功能——导致关键token被驱逐——从而降低LLM性能。为解决该问题——我们提出CompressKV——面向GQA式LLM的资源高效KV缓存压缩框架。CompressKV不从所有头聚合注意力分数——而是识别同时捕捉提示首尾token与语义重要中段证据的语义检索头(SRH)——并用它们选择应保留KV对的token。此外——CompressKV按层驱逐误差的离线估计跨层分配缓存预算。LongBench与大海捞针实验表明CompressKV跨内存预算持续超越既有KV缓存驱逐方法。值得注意的是——它在LongBench问答任务上仅用3%的KV缓存保持全缓存性能的97%以上——在大海捞针上仅用0.7%的KV存储达到90%准确率。这些结果展示长上下文LLM推理更优的资源-性能权衡。代码公开于: https://github.com/TUDa-HWAI/CompressKV。

CompressKV:语义检索引导的KV缓存压缩实现资源高效长上下文LLM推理:论文配图
图 4:与基准方法相比,不同 KV 缓存预算下 16 个 LongBench 数据集的平均性能。