arXiv:2411.18077cs.CLcs.LG2024-11被引 6

用2比特分层压缩注意力缓存,大幅降低大模型推理内存占用。

MiniKV: Pushing the Limits of LLM Inference via 2-Bit Layer-Discriminative KV Cache

  • 采用2比特分层判别式缓存策略,精准压缩注意力键值对。
  • 在长文本任务中实现86%缓存压缩率,准确率恢复超98.5%。
  • 适配FlashAttention,适合高并发长序列大模型部署场景。

由于大语言模型(LLM)的内存与计算需求巨大,其实际服务面临严峻挑战。本文聚焦于KV缓存这一关键瓶颈——它在长上下文任务中占据大量内存。为此,提出MiniKV,一种新型2比特层判别式KV缓存优化方法,在保持长上下文任务精度的同时显著压缩缓存空间。更重要的是,开发了专用CUDA内核,使MiniKV兼容FlashAttention。在多种长上下文任务上的实验表明,MiniKV实现了86%的缓存压缩比,同时准确率恢复超过98.5%,优于现有最先进方法,并带来显著的系统性能提升。

原文摘要 · Abstract (English)

How to efficiently serve LLMs in practice has become exceptionally challenging due to their prohibitive memory and computation requirements. In this study, we investigate optimizing the KV cache, whose memory footprint poses a critical bottleneck in LLM inference, especially when dealing with long context tasks. To tackle the challenge, we introduce MiniKV, a KV cache optimization method that simultaneously preserves long context task accuracy while significantly reducing KV cache size via a novel 2-bit layer-discriminative KV cache. More importantly, we develop specialized CUDA kernels to make MiniKV compatible with FlashAttention. Experiments on a wide range of long context tasks show that MiniKV effectively achieves 86% KV cache compression ratio while recovering over 98.5% of accuracy, outperforming state-of-the-art methods while achieving excellent measured system performance improvements.

KV缓存大模型推理内存优化2比特量化

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