针对长文本生成中缓存内存过大的问题,提出按层定制缓存大小的方法。
XKV: Personalized KV Cache Memory Reduction for Long-Context LLM Inference
- 按层个性化分配缓存空间,而非统一配置。
- 平均降低61.6%缓存内存,吞吐提升最高达5.5倍。
- 适合需要高效长文本推理的场景,如文档摘要、对话系统。
生成式大语言模型在推理时逐个生成输出标记,导致大量冗余计算。广泛采用的键值缓存(KV-Cache)机制在时间与空间复杂度间权衡,但缓存数据随上下文增长而持续占用内存,易超出现代加速器(如GPU)的有限内存容量,尤其在长上下文任务中更为严重。现有方法通过淘汰影响较小的缓存数据来减少内存消耗,但由于各网络层缓存分配静态,实际效果不佳。本文发现不同层的缓存数据对精度影响差异显著,并通过实验与理论验证了这一点。进一步提出对每层进行个性化缓存大小配置,可大幅减少内存且保持相近精度。将缓存分配建模为组合优化问题,给出全局最优解。设计轻量级模型的小规模采样推理,快速捕捉层间差异并输入个性化算法。在真实数据集上的大量实验表明,该方法平均降低61.6%的KV缓存内存,计算效率提升2.1倍,吞吐最高提升5.5倍。
原文摘要 · Abstract (English)
Recently the generative Large Language Model (LLM) has achieved remarkable success in numerous applications. Notably its inference generates output tokens one-by-one, leading to many redundant computations. The widely-used KV-Cache framework makes a compromise between time and space complexities. However, caching data generates the increasingly growing memory demand, that can quickly exhaust the limited memory capacity of the modern accelerator like GPUs, particularly in long-context inference tasks. Existing studies reduce memory consumption by evicting some of cached data that have less important impact on inference accuracy. But the benefit in practice is far from ideal due to the static cache allocation across different LLM network layers. This paper observes that the layer-specific cached data have very different impacts on accuracy. We quantify this difference, and give experimental and theoretical validation. We accordingly make a formal analysis and shows that customizing the cache size for each layer in a personalized manner can yield a significant memory reduction, while still providing comparable accuracy. We simulate the cache allocation as a combinatorial optimization problem and give a global optimal solution. In particular, we devise a mini- and sampling-based inference over a lightweight variant of the LLM model, so as to quickly capture the difference and then feed it into the personalized algorithms. Extensive experiments on real-world datasets demonstrate that our proposals can reduce KV cache memory consumption by 61.6% on average, improve computational efficiency by 2.1x and then increase the throughput by up to 5.5x.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。