AnTKV: Anchor Token-Aware Sub-Bit Vector Quantization for KV Cache in Large Language Models
This work addresses memory efficiency for deploying large language models, offering a novel method that improves performance in ultra-low-bit regimes, though it is incremental in the context of existing quantization techniques.
The paper tackles the problem of reducing memory footprint in Large Language Models' KV cache through ultra-low-bit quantization, proposing AnTKV, which uses anchor token-aware vector quantization to mitigate accuracy degradation, achieving a perplexity of 6.32 at 1-bit on Mistral-7B and enabling LLaMA3-8B to scale to 840K tokens with 3.5x higher decoding throughput.
Quantization has emerged as an effective and lightweight solution to reduce the memory footprint of the KV cache in Large Language Models. Nevertheless, minimizing the accuracy degradation caused by ultra-low-bit KV cache quantization remains a significant challenge. While scalar quantization is constrained by 1-bit bound, vector quantization exploits intra-vector correlations and enables sub-bit regimes, making it more suitable for ultra-low-bit quantization. To further mitigate quantization-induced degradation, we reveal that the degradation is highly uneven across tokens in attention quality. To investigate this unevenness, we introduce anchor score to measure each token's sensitivity to quantization. Our analysis and experiments show that preserving a small subset (1\%) of tokens with the highest Anchor Score significantly mitigates accuracy loss under aggressive quantization. We propose AnTKV, a dual-stage framework that leverages anchor token-aware vector quantization to compress the KV cache. It combines offline token-aware centroids learning and online anchor token selection to balance compression and accuracy. To enable efficient deployment, we design an online anchor token selection kernel compatible with FlashAttention. It allows LLaMA3-8B to scale to 840K tokens on a single 80GB A100, while delivering up to $3.5\times$ higher decoding throughput over the FP16 baseline. Experiments demonstrate that AnTKV matches or surpasses prior methods at 4-bit, and significantly reduce perplexity under ultra-low-bit quantization, achieving 6.32 at 1-bit on Mistral-7B, compared to 7.25 for CQ and 15.36 for KVQuant.