LaCache: Ladder-Shaped KV Caching for Efficient Long-Context Modeling of Large Language Models

文章主要内容和创新点

主要内容

本文针对大型语言模型（LLMs）在长上下文处理中面临的键值（KV）缓存内存瓶颈问题，提出了一种无需训练的KV缓存优化框架LaCache。该框架旨在平衡LLMs的长程建模能力与连续生成时的内存效率，解决现有方法在长序列处理中要么内存不足（OOM）、要么长程依赖捕捉能力弱的问题。

通过设计阶梯状KV缓存模式和迭代压缩机制，LaCache在固定内存预算下扩展了长程依赖捕捉的范围，同时支持无限长序列的连续生成。实验验证显示，在多个模型（如Llama2、Llama3、SmolLM2等）和任务（语言建模、长上下文理解等）中，LaCache的性能均优于现有方法（如StreamingLLM、H2O等），且与高效注意力实现（如FlashAttention）兼容，兼顾精度与吞吐量。

创新点

1. 阶梯状KV缓存模式：
   不同于现有方法在所有层存储相同的KV对，LaCache在浅层保留早期token的KV状态，在深层逐渐聚焦于后期token，形成阶梯状结构。这一设计在固定内存预算下扩展了长程依赖的捕捉范围，提升了长程建模能力。

2. 迭代压缩机制：
   当KV缓存达到容量上限时，通过阶梯状模式对已压缩的旧缓存进行进一步压缩，释放空间供新token使用。该机制对旧token压缩更激进、对新token压缩更温和，在保证内存效率的同时优先保留近期信息，支持无限长序列的连续生成。

3. 兼容性与效率