XQuant: Superare il Muro della Memoria per l'Inferenza di LLM con la Rimaterializzazione della Cache KV

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

August 14, 2025

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

cs.AI

Abstract

Sebbene l'inferenza dei modelli linguistici di grandi dimensioni (LLM) sia emersa come un carico di lavoro critico per molte applicazioni downstream, inferire efficientemente gli LLM è impegnativo a causa dell'ingente impronta di memoria e dei requisiti di larghezza di banda. Parallelamente, le capacità di calcolo hanno superato costantemente sia la capacità di memoria che la larghezza di banda negli ultimi decenni, una tendenza che rimane evidente nell'hardware GPU moderno e che aggrava la sfida dell'inferenza degli LLM. Di conseguenza, stanno emergendo nuovi algoritmi che scambiano un aumento del calcolo con una riduzione delle operazioni di memoria. A tal fine, presentiamo XQuant, che sfrutta questa tendenza, consentendo una riduzione di un ordine di grandezza nel consumo di memoria attraverso la quantizzazione a basso bit, con sostanziali vantaggi in termini di accuratezza rispetto ai metodi di quantizzazione della cache KV allo stato dell'arte. Raggiungiamo questo obiettivo quantizzando e memorizzando nella cache le attivazioni di input X, invece di utilizzare la standard cache KV, e poi rimaterializzando le Chiavi e i Valori al volo durante l'inferenza. Ciò si traduce in un immediato risparmio di memoria di 2 volte rispetto alla cache KV. Applicando XQuant, otteniamo un risparmio di memoria fino a ~7,7 volte con una degradazione della perplessità <0,1 rispetto alla baseline FP16. Inoltre, il nostro approccio sfrutta il fatto che i valori X sono simili tra i livelli. Basandoci su questa osservazione, introduciamo XQuant-CL, che sfrutta la similarità cross-layer negli embedding X per una compressione estrema. Su diversi modelli, XQuant-CL raggiunge un risparmio di memoria fino a 10 volte rispetto alla baseline FP16 con solo una degradazione della perplessità di 0,01, e un risparmio di memoria di 12,5 volte con solo una degradazione della perplessità di 0,1. XQuant sfrutta le capacità di calcolo in rapida crescita delle piattaforme hardware per eliminare il collo di bottiglia della memoria, superando i metodi di quantizzazione della cache KV allo stato dell'arte e raggiungendo un'accuratezza quasi-FP16 su un'ampia gamma di modelli.

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.