Table of Contents
Fetching ...

PackKV: Reducing KV Cache Memory Footprint through LLM-Aware Lossy Compression

Bo Jiang, Taolue Yang, Youyuan Liu, Xubin He, Sheng Di, Sian Jin

TL;DR

PackKV addresses the KV cache memory bottleneck in long-context LLM inference by introducing a novel, LLM-aware lossy compression framework that fuses quantization with ultra-light lossless encoding and a computation-aware decompression kernel. The method exploits permutation invariance in attention to reorder KV vectors for better bit-packing while keeping the K and V pairs bound together, and it integrates decompression directly into matrix–vector multiplication to eliminate memory writebacks. Across six models and six benchmarks, PackKV achieves on average 153.2% higher memory reduction for K and 179.6% for V compared with state-of-the-art quantization baselines, while delivering substantial throughput gains—75.7% for K and 171.7% for V—on A100 and RTX Pro 6000 GPUs. The approach is orthogonal to pruning or CPU–GPU offloading, demonstrates scalable performance across GPUs, and paves the way for longer-context inference with limited hardware resources.

Abstract

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

PackKV: Reducing KV Cache Memory Footprint through LLM-Aware Lossy Compression

TL;DR

PackKV addresses the KV cache memory bottleneck in long-context LLM inference by introducing a novel, LLM-aware lossy compression framework that fuses quantization with ultra-light lossless encoding and a computation-aware decompression kernel. The method exploits permutation invariance in attention to reorder KV vectors for better bit-packing while keeping the K and V pairs bound together, and it integrates decompression directly into matrix–vector multiplication to eliminate memory writebacks. Across six models and six benchmarks, PackKV achieves on average 153.2% higher memory reduction for K and 179.6% for V compared with state-of-the-art quantization baselines, while delivering substantial throughput gains—75.7% for K and 171.7% for V—on A100 and RTX Pro 6000 GPUs. The approach is orthogonal to pruning or CPU–GPU offloading, demonstrates scalable performance across GPUs, and paves the way for longer-context inference with limited hardware resources.

Abstract

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
Paper Structure (26 sections, 4 equations, 17 figures, 5 tables, 1 algorithm)

This paper contains 26 sections, 4 equations, 17 figures, 5 tables, 1 algorithm.

Figures (17)

  • Figure 1: KV Cache behavior during LLM Inference, $M_{Q,K,V}$ is the mapping matrix to $K,Q,V$ vector. Each $Q$ vector performs dot product with every $K$ vector to generate a weight after softmax operation for each $V$ vector. All the $V$ vectors multiplied with their weights and aggregate to Attn Output.
  • Figure 2: Overview of PackKV. Compared with vanilla and previous quantization solutions. (1), We rearrange the order of quantized kv cache vector to benefit the following lossless bit-packing encode. (2). Bit-packing the quantized kv cache integers in channel direction to achieve extra memory saving. And to eliminate the decompress overhead we developed a Computation-aware Decompression + matrix-vector Multiplication method.
  • Figure 3: Histogram of K and V quantized integer tensor, x-axis is quantized integer value, y-axis is the frequency of the value.
  • Figure 4: A visualization of a sampled K after quantization to integer, horizontal direction is hidden_dim and vertical is context_len (channel) direction.
  • Figure 5: KV quantization granularities based on KV cache dimensions: token-wise and channel-wise quantization.
  • ...and 12 more figures