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

TL;DR

MiniKV is a novel 2-bit layer-discriminative KV cache method that significantly reduces memory usage in LLM inference while maintaining high accuracy, enabling more efficient long context processing.

Contribution

We introduce MiniKV, a 2-bit layer-discriminative KV cache technique with specialized CUDA kernels, improving memory efficiency and accuracy in LLM inference for long context tasks.

Findings

86% KV cache compression ratio achieved

Over 98.5% accuracy recovery

Outperforms state-of-the-art methods in efficiency

Abstract

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.