Hessian-Aware KV Cache Quantization for LLMs

As large language models continue to advance, handling extended context lengths becomes crucial for improving natural language processing tasks. However, it also introduces significant computational challenges, particularly in managing KV cache memory. To address this, we present a novel Hessian-aware KV cache quantization method for LLMs. This method employs a weight-aware head-wise Hessian analysis and a three-tiered mixed-precision quantization strategy. It assigns high precision to magnitude-based outliers, intermediate precision to Hessian-based important elements, and low precision to the rest, achieving an optimal balance between memory compression and accuracy. Our extensive testing across various LLM configurations and datasets shows that this method achieves significant KV cache compression rates, ranging from 8.21x to 9.7x while maintaining minimal deviation in accuracy compared to the FP16 baseline. This enhancement enables a remarkable 3.8x higher throughput (tokens/second) than the FP16 baseline, highlighting the practical utility of our approach in addressing computational challenges in LLMs.