XQuant: Superando a Barreira de Memória na Inferência de LLMs com Rematerialização do Cache KV

XQuant: Breaking the Memory Wall for LLM Inference with KV Cache Rematerialization

August 14, 2025

Autores: Aditya Tomar, Coleman Hooper, Minjae Lee, Haocheng Xi, Rishabh Tiwari, Wonjun Kang, Luca Manolache, Michael W. Mahoney, Kurt Keutzer, Amir Gholami

cs.AI

Resumo

Embora a inferência de LLMs tenha surgido como uma carga de trabalho crítica para muitas aplicações downstream, inferir LLMs de forma eficiente é desafiador devido à pegada de memória substencial e aos requisitos de largura de banda. Paralelamente, as capacidades de computação têm superado consistentemente tanto a capacidade de memória quanto a largura de banda nas últimas décadas, uma tendência que permanece evidente no hardware moderno de GPUs e exacerba o desafio da inferência de LLMs. Como resultado, novos algoritmos estão surgindo que trocam o aumento da computação pela redução das operações de memória. Para esse fim, apresentamos o XQuant, que aproveita essa tendência, permitindo uma redução de uma ordem de magnitude no consumo de memória por meio de quantização de baixo bit, com benefícios substanciais de precisão em relação aos métodos de quantização de cache KV state-of-the-art. Isso é alcançado quantizando e armazenando em cache as ativações de entrada da camada X, em vez de usar o cache KV padrão, e então rematerializando as Chaves e Valores dinamicamente durante a inferência. Isso resulta em uma economia imediata de 2 vezes na memória em comparação com o cache KV. Ao aplicar o XQuant, alcançamos uma economia de memória de até ~7,7 vezes com uma degradação de perplexidade <0,1 em comparação com a linha de base FP16. Além disso, nossa abordagem aproveita o fato de que os valores de X são semelhantes entre as camadas. Com base nessa observação, introduzimos o XQuant-CL, que explora a similaridade entre camadas nas incorporações X para compressão extrema. Em diferentes modelos, o XQuant-CL atinge uma economia de memória de até 10 vezes em relação à linha de base FP16 com apenas 0,01 de degradação de perplexidade, e 12,5 vezes de economia de memória com apenas 0,1 de degradação de perplexidade. O XQuant aproveita o rápido aumento das capacidades de computação das plataformas de hardware para eliminar o gargalo de memória, superando os métodos de quantização de cache KV state-of-the-art e alcançando precisão próxima à FP16 em uma ampla gama de modelos.

English

Although LLM inference has emerged as a critical workload for many downstream applications, efficiently inferring LLMs is challenging due to the substantial memory footprint and bandwidth requirements. In parallel, compute capabilities have steadily outpaced both memory capacity and bandwidth over the last few decades, a trend that remains evident in modern GPU hardware and exacerbates the challenge of LLM inference. As such, new algorithms are emerging that trade increased computation for reduced memory operations. To that end, we present XQuant, which takes advantage of this trend, enabling an order-of-magnitude reduction in memory consumption through low-bit quantization with substantial accuracy benefits relative to state-of-the-art KV cache quantization methods. We accomplish this by quantizing and caching the layer input activations X, instead of using standard KV caching, and then rematerializing the Keys and Values on-the-fly during inference. This results in an immediate 2times memory savings compared to KV caching. By applying XQuant, we achieve up to sim 7.7times memory savings with <0.1 perplexity degradation compared to the FP16 baseline. Furthermore, our approach leverages the fact that X values are similar across layers. Building on this observation, we introduce XQuant-CL, which exploits the cross-layer similarity in the X embeddings for extreme compression. Across different models, XQuant-CL attains up to 10times memory savings relative to the FP16 baseline with only 0.01 perplexity degradation, and 12.5times memory savings with only 0.1 perplexity degradation. XQuant exploits the rapidly increasing compute capabilities of hardware platforms to eliminate the memory bottleneck, while surpassing state-of-the-art KV cache quantization methods and achieving near-FP16 accuracy across a wide range of models.