XQuant: Superando la Barrera de Memoria para la Inferencia de LLM con Rematerialización de la Caché 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

Resumen

Aunque la inferencia de LLM ha surgido como una carga de trabajo crítica para muchas aplicaciones posteriores, inferir LLMs de manera eficiente es un desafío debido a la considerable huella de memoria y los requisitos de ancho de banda. En paralelo, las capacidades de cómputo han superado constantemente tanto la capacidad de memoria como el ancho de banda durante las últimas décadas, una tendencia que sigue siendo evidente en el hardware moderno de GPU y que exacerba el desafío de la inferencia de LLM. Como resultado, están surgiendo nuevos algoritmos que intercambian un mayor cómputo por una reducción en las operaciones de memoria. Con este fin, presentamos XQuant, que aprovecha esta tendencia, permitiendo una reducción de un orden de magnitud en el consumo de memoria mediante la cuantización de bajo bit, con beneficios sustanciales en precisión en comparación con los métodos de cuantización de caché KV más avanzados. Logramos esto cuantizando y almacenando en caché las activaciones de entrada de la capa X, en lugar de utilizar el almacenamiento en caché KV estándar, y luego rematerializando las Claves y Valores sobre la marcha durante la inferencia. Esto resulta en un ahorro inmediato de memoria de 2 veces en comparación con el almacenamiento en caché KV. Al aplicar XQuant, logramos un ahorro de memoria de hasta ~7.7 veces con una degradación de <0.1 en la perplejidad en comparación con la línea base FP16. Además, nuestro enfoque aprovecha el hecho de que los valores de X son similares entre capas. Basándonos en esta observación, introducimos XQuant-CL, que explota la similitud entre capas en los embeddings de X para una compresión extrema. En diferentes modelos, XQuant-CL alcanza un ahorro de memoria de hasta 10 veces en comparación con la línea base FP16 con solo una degradación de 0.01 en la perplejidad, y un ahorro de memoria de 12.5 veces con solo una degradación de 0.1 en la perplejidad. XQuant aprovecha las capacidades de cómputo en rápido aumento de las plataformas de hardware para eliminar el cuello de botella de la memoria, superando los métodos de cuantización de caché KV más avanzados y logrando una precisión cercana a FP16 en una amplia 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.