通过梯度分析动态分配关键层精度,显著降低LLM推理内存占用。
KVmix: Gradient-Based Layer Importance-Aware Mixed-Precision Quantization for KV Cache
- 基于梯度重要性分析,为不同层分配不同精度
- 实现2.19比特键、2.38比特值的极低配置,内存压缩4.9倍
- 适合资源受限场景下高精度长文本生成
大型语言模型(LLMs)推理过程中,键值(KV)缓存的高内存需求严重限制了其在资源受限平台上的部署。量化可有效缓解KV缓存带来的内存压力。然而,现有方法或依赖静态统一精度分配,或无法在长上下文任务中动态优先处理关键KV,导致必须在内存、精度与吞吐间权衡。本文提出一种新型混合精度量化方法KVmix,利用梯度重要性分析评估各键值投影矩阵对模型损失的影响,实现层级粒度的比特位分配。该方法动态提升关键层精度,对影响小的层进行激进量化,实现精度与效率的可调平衡。同时引入动态长上下文优化策略,自适应保留近期关键令牌的全精度KV对,压缩旧数据,以低内存开销实现高质量序列生成。此外,KVmix提供高效的低比特量化及CUDA内核,优化计算开销。在Llama和Mistral等LLM上,KVmix以极低量化配置(键2.19比特,值2.38比特)实现近无损推理性能,内存压缩达4.9倍,推理吞吐提升5.3倍。
原文摘要 · Abstract (English)
The high memory demands of the Key-Value (KV) Cache during the inference of Large Language Models (LLMs) severely restrict their deployment in resource-constrained platforms. Quantization can effectively alleviate the memory pressure caused by KV Cache. However, existing methods either rely on static one-size-fits-all precision allocation or fail to dynamically prioritize critical KV in long-context tasks, forcing memory-accuracy-throughput tradeoffs. In this work, we propose a novel mixed-precision quantization method for KV Cache named KVmix. KVmix leverages gradient-based importance analysis to evaluate how individual Key and Value projection matrices affect the model loss, enabling layer-specific bit-width allocation for mix-precision quantization. It dynamically prioritizes higher precision for important layers while aggressively quantizing less influential ones, achieving a tunable balance between accuracy and efficiency. KVmix also introduces a dynamic long-context optimization strategy that adaptively keeps full-precision KV pairs for recent pivotal tokens and compresses older ones, achieving high-quality sequence generation with low memory usage. Additionally, KVmix provides efficient low-bit quantization and CUDA kernels to optimize computational overhead. On LLMs such as Llama and Mistral, KVmix achieves near-lossless inference performance with extremely low quantization configuration (Key 2.19bit Value 2.38bit), while delivering a remarkable 4.9x memory compression and a 5.3x speedup in inference throughput.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。