Kitty: A 2-bit KV Cache Quantization Scheme for Large Model Inference

The Kitty project proposes an innovative KV cache quantization method that achieves high-precision large language model (LLM) inference acceleration using only 2-bit quantization, leveraging dynamic channel precision enhancement technology. This scheme aims to address the bottleneck of rapidly expanding GPU memory usage by KV cache in LLM inference, and has significant practical deployment value.

In the inference process of large language models (LLMs), KV cache is a core component of the attention mechanism, storing historical token key-value pairs. As model size and sequence length increase, the GPU memory usage of KV cache expands rapidly, becoming the main bottleneck for long-context inference. Traditional quantization methods adopt a uniform precision strategy, ignoring the differential sensitivity of channels, leading to severe precision loss at extremely low bits (e.g., 2-bit), making practical deployment difficult.

The core of Kitty is the dynamic channel precision enhancement algorithm. The research team found: 1) There are significant differences in the numerical distribution of different channels and their contributions to attention scores; 2) Different channels have different sensitivities to quantization errors. Based on this, the algorithm adopts strategies such as adaptive precision allocation (higher precision for key channels), lightweight runtime decision-making (low-overhead online analysis), and compatibility with the 2-bit baseline (maintaining low storage while restoring performance).

Kitty implements an algorithm-system co-design: 1) Plug-and-play integration: Modified the HuggingFace Transformers library to seamlessly compatible with existing models and workflows; 2) Apptainer containerized deployment: Supports CUDA 12.1 environment, writable overlay images, and multi-instance read-only mode; 3) Complete evaluation framework: Integrated lm-evaluation-harness, supporting latency benchmarking, precision simulation verification, and multi-task evaluation.

The application value of Kitty is significant: 1) Memory efficiency improvement: 2-bit quantization reduces KV cache usage to 1/8 of FP16, allowing a single card to handle longer contexts; 2) Inference throughput optimization: Reduced memory usage allows larger batch sizes, improving server concurrency; 3) Edge device feasibility: Extremely low memory requirements promote the application of LLMs on resource-constrained devices.

Currently, Kitty is in the peer review stage of academic conferences and relies on a specific Transformers version (hf-4.53.2 branch); compatibility with upstream versions needs continuous attention. Future directions include: supporting more model architectures (e.g., Mamba, RWKV), exploring the feasibility of 1-bit quantization, and combining acceleration technologies such as speculative decoding.

The Kitty project, through in-depth understanding of the characteristics of the attention mechanism, achieves a breakthrough in maintaining high performance at extremely low precision, demonstrating the potential of algorithm innovation in solving large model deployment challenges, and paving the way for the popular application of long-context large models.