用压缩技术大幅减少大模型推理时的显存占用,提升速度和效率。
PackKV: Reducing KV Cache Memory Footprint through LLM-Aware Lossy Compression
- 针对KV缓存特性设计专用有损压缩算法,兼顾压缩率与计算效率。
- 相比顶尖量化方法,K/V缓存内存减少超150%,推理吞吐提升75%以上。
- 适用于长文本生成场景,适合需要高效推理的部署者使用。
基于Transformer的大语言模型在诸多实际应用中展现出巨大潜力,但长序列推理仍面临显著挑战,主要源于键值(KV)缓存带来的巨大内存开销,其大小随序列长度和批量增大可达到数GB。本文提出一种通用且高效的KV缓存管理框架PackKV,专为长上下文生成优化。PackKV引入针对KV缓存数据特性的新型有损压缩技术,通过压缩算法与系统架构的协同设计,在保持高计算效率的同时,兼容缓存动态增长的特性。实验表明,在与当前最先进量化方法同等甚至更低的精度损失下,PackKV平均实现K缓存153.2%、V缓存179.6%的内存压缩率。此外,其执行吞吐极高,有效消除解压开销,并加速矩阵-向量乘法操作。在A100和RTX Pro 6000 GPU上,相比cuBLAS矩阵-向量乘法核,平均吞吐分别提升75.7%(K)和171.7%(V),同时降低显存带宽需求。代码已开源:https://github.com/BoJiang03/PackKV。
原文摘要 · Abstract (English)
Transformer-based large language models (LLMs) have demonstrated remarkable potential across a wide range of practical applications. However, long-context inference remains a significant challenge due to the substantial memory requirements of the key-value (KV) cache, which can scale to several gigabytes as sequence length and batch size increase. In this paper, we present \textbf{PackKV}, a generic and efficient KV cache management framework optimized for long-context generation. %, which synergistically supports both latency-critical and throughput-critical inference scenarios. PackKV introduces novel lossy compression techniques specifically tailored to the characteristics of KV cache data, featuring a careful co-design of compression algorithms and system architecture. Our approach is compatible with the dynamically growing nature of the KV cache while preserving high computational efficiency. Experimental results show that, under the same and minimum accuracy drop as state-of-the-art quantization methods, PackKV achieves, on average, \textbf{153.2}\% higher memory reduction rate for the K cache and \textbf{179.6}\% for the V cache. Furthermore, PackKV delivers extremely high execution throughput, effectively eliminating decompression overhead and accelerating the matrix-vector multiplication operation. Specifically, PackKV achieves an average throughput improvement of \textbf{75.7}\% for K and \textbf{171.7}\% for V across A100 and RTX Pro 6000 GPUs, compared to cuBLAS matrix-vector multiplication kernels, while demanding less GPU memory bandwidth. Code available on https://github.com/BoJiang03/PackKV
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。