arXiv:2505.19586cs.CL2025-05ACL被引 12

TailorKV通过混合压缩提升长文本推理,让单张显卡跑128k上下文。

TailorKV: A Hybrid Framework for Long-Context Inference via Tailored KV Cache Optimization

  • 按层区分缓存需求,对关键层保留完整信息,对其他层做量化与卸载
  • 在128k上下文中实现近乎无损性能,每轮生成仅需82毫秒
  • 适合需要超长上下文的部署场景,如法律、医疗文档分析

生成式大语言模型中的键值(KV)缓存带来显著内存开销。现有方法通过缓存卸载或压缩缓解该问题,但完全加载缓存受CPU-GPU通信带宽限制导致延迟高,而激进压缩则造成性能明显下降。我们发现,模型某些层需保持全局信息,不适合选择性加载;而其他层主要关注少数激活强度高的令牌,存在较大量化误差风险。基于此观察,提出互补策略:加载主导令牌并量化所有令牌。据此设计混合压缩方案TailorKV,构建兼具高效性与硬件友好性的推理框架。大量长上下文实验表明,尾随压缩设置下几乎无性能损失,优于当前最优方法。尤其在128k上下文的Llama-3.1-8B模型中,可在单张RTX 3090 GPU上运行,解码时达到每标记82毫秒的吞吐。

原文摘要 · Abstract (English)

The Key-Value (KV) cache in generative large language models (LLMs) introduces substantial memory overhead. Existing works mitigate this burden by offloading or compressing the KV cache. However, loading the entire cache incurs significant latency due to PCIe bandwidth bottlenecks in CPU-GPU communication, while aggressive compression causes notable performance degradation. We identify that certain layers in the LLM need to maintain global information and are unsuitable for selective loading. In contrast, other layers primarily focus on a few tokens with dominant activations that potentially incur substantial quantization error. This observation leads to a key insight that loading dominant tokens and quantizing all tokens can complement each other. Building on this insight, we propose a hybrid compression method, TailorKV, which seamlessly integrates quantization and offloading. TailorKV develops an inference framework along with a hardware-friendly implementation that leverages these complementary characteristics. Extensive long-context evaluations exhibit that TailorKV achieves nearly lossless performance under aggressive compression settings, outperforming the state-of-the-art. Particularly, the Llama-3.1-8B with 128k context can be served within a single RTX 3090 GPU, reaching 82 ms per token during decoding.

长上下文缓存优化混合压缩

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。