论文洞察：基于重要性感知的多层级前缀KV Cache存储系统

研究背景

本文基于浙江大学、华为在FAST25上发表的研究成果《IMPRESS: An Importance-Informed Multi-Tier Prefix KV Storage System for Large Language Model Inference》进行深入分析。

当CPU内存不足时需要将前缀KV Cache存储在磁盘上，由于磁盘I/O延迟较高，LLM推理时复用KV Cache可能导致TTFT时延增加。本文提出IMPRESS，一种基于重要性感知的多层级前缀KV Cache存储系统，通过仅加载重要的前缀KV Cache来减少LLM推理的I/O延迟。总体而言，IMPRESS属于KV Cache加载方面的LLM推理优化技术。

核心问题识别

LLM推理时通常在用户查询前添加长上下文以提高输出质量，并且这些上下文在多个查询中可能存在部分或完全重复的情况，因此现有系统通过存储和复用这些上下文的KV Cache来减少冗余计算并降低TTFT时间。然而，当CPU内存不足时，KV Cache需要存储在磁盘上，由于磁盘I/O延迟高，复用磁盘中的KV Cache在极大程度上会增加TTFT时延。

1.存储瓶颈：因GPU/CPU存在内存容量瓶颈，长序列或大批量场景下需将前缀KV Cache存储到磁盘，而磁盘I/O延迟成为新瓶颈，占TTFT的51%-98%；

2.低效识别：现有重要KV Cache识别需加载全部前缀KV Cache到GPU显存上来计算注意力权重，I/O开销大；

3.存储与缓存优化不足：传统系统将连续KV Cache合并为块（chunk），读取重要KV Cache时会附带无关数据，且基于块访问模式的缓存管理未考虑KV Cache重要性，导致缓存命中率低。

核心设计与创新

IMPRESS是一种基于重要性感知的多层级前缀KV Cache存储系统，整合GPU内存、CPU内存和磁盘三层存储，通过有选择的加载磁盘中重要前缀KV Cache以减少I/O延迟，从而提高LLM推理性能。

1.相似性引导的重要KVCache识别（ITF）

• 由于同一Transformer层中不同注意力头的重要token索引集高度相似，因此随机选取所有注意力头中的3个作为探测头（probe heads），仅加载KV Cache中K值到GPU显存来计算注意力权重，通过相似度阈值生成重要token索引集，从而避免加载全部注意力头的K值；
  2.基于重要性感知的KVCache管理
• KVCache重排序：按token重要性定期重组磁盘上的KV Cache块，提高块内重要KV Cache密度，降低读取重要KV Cache时的无效数据的加载频率。同时兼容基于基数树的前缀检索，不破坏现有元数据结构；
• 基于Score的KVCache管理：为每个块计算“访问频率×重要KV Cache比例”的Score，优先将高Score的KV Cache缓存到GPU内存，采用最小堆管理GPU/CPU缓存，避免缓存冗余，以提高重要KV Cache的GPU缓存命中率。

图IMPRESS系统

实验验证

基于FlexGen实现了IMPRESS，并在OPT-6.7B、OPT-13B和OPT-30B等模型上针对IMPRESS、ReComp、AS-like、AS+H2O+LRU、AS+H2O+LFU算法进行了测试对比。

• 相比ReComp、AS-like、AS+H2O+LRU、AS+H2O+LFU算法，IMPRESS的TTFT性能提升1.2-2.8倍，KV Cache的I/O加载开销减少1.5-3.8倍。
• IMPRESS为每个KV chunk添加一个score得分，其空间开销＜0.5%，对整体系统的影响可忽略。

图在四个数据集下，不同算法的TTFT时间对比

图不同算法的I/O开销对比

研究价值

IMPRESS有效解决了LLM长前缀推理中的磁盘I/O瓶颈问题，为长上下文LLM推理服务提供了新的存储优化方案，且重要token识别过程开销很小，所以对现有LLM推理框架的影响较小。

论文链接

https://www.usenix.org/conference/fast25/presentation/chen-weijian-impress

以上内容转载自微信公众号—大模型软硬协同优化，链接：https://mp.weixin.qq.com/s/nXlQ