针对长文本推理中键值缓存内存占用过高的问题，提出按重要性分配内存的新方法。
BaKlaVa -- Budgeted Allocation of KV cache for Long-context Inference
- 根据缓存重要性动态分配内存，而非统一处理所有注意力头
- 在保持基线性能前提下实现最高70%的压缩率
- 适合需要高效长序列推理的部署场景
在大语言模型推理中，键值缓存（KV-caches）对降低时间复杂度至关重要，但会随上下文长度线性增加显存占用。现有工作虽探索了缓存淘汰与压缩策略，但通常对所有注意力头采用统一缓存管理，导致性能不佳。本文提出BaKlaVa，通过一次性的性能分析，评估每个KV缓存的重要性，并为各缓存分配最优内存预算。实验基于LLaMA-3-8B和Qwen2.5-7B模型，结果表明：在保持原始性能的前提下，最高可实现70%的压缩率；在更高压缩比下，准确率提升可达一个数量级。
原文摘要 · Abstract (English)
In Large Language Model (LLM) inference, Key-Value (KV) caches (KV-caches) are essential for reducing time complexity. However, they result in a linear increase in GPU memory as the context length grows. While recent work explores KV-cache eviction and compression policies to reduce memory usage, they often consider uniform KV-caches across all attention heads, leading to suboptimal performance. We introduce BaKlaVa, a method to allocate optimal memory for individual KV-caches across the model by estimating the importance of each KV-cache. Our empirical analysis demonstrates that not all KV-caches are equally critical for LLM performance. Using a one-time profiling approach, BaKlaVa assigns optimal memory budgets to each KV-cache. We evaluated our method on LLaMA-3-8B, and Qwen2.5-7B models, achieving up to a 70\% compression ratio while keeping baseline performance and delivering up to an order-of-magnitude accuracy improvement at higher compression levels.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。