XQuant: Преодоление ограничений памяти при выводе больших языковых моделей за счет рематериализации кэша ключей и значений

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

August 14, 2025

Авторы: Aditya Tomar, Coleman Hooper, Minjae Lee, Haocheng Xi, Rishabh Tiwari, Wonjun Kang, Luca Manolache, Michael W. Mahoney, Kurt Keutzer, Amir Gholami

cs.AI

Аннотация

Хотя вывод больших языковых моделей (LLM) стал критически важной задачей для многих приложений, эффективный вывод LLM является сложной задачей из-за значительного объема памяти и требований к пропускной способности. В то же время вычислительные возможности на протяжении последних десятилетий стабильно опережали как объем памяти, так и пропускную способность, и эта тенденция остается очевидной в современных GPU, что усугубляет сложность вывода LLM. В связи с этим появляются новые алгоритмы, которые жертвуют увеличением вычислений ради сокращения операций с памятью. В этом контексте мы представляем XQuant, который использует эту тенденцию, обеспечивая снижение потребления памяти на порядок за счет низкобитового квантования с существенным улучшением точности по сравнению с современными методами квантования кэша ключей и значений (KV). Мы достигаем этого, квантуя и кэшируя входные активации слоев X вместо использования стандартного кэширования KV, а затем динамически восстанавливая ключи и значения во время вывода. Это приводит к немедленному сокращению памяти в 2 раза по сравнению с кэшированием KV. Применяя XQuant, мы достигаем сокращения памяти до ~7.7 раз с ухудшением перплексии менее чем на 0.1 по сравнению с базовым уровнем FP16. Кроме того, наш подход использует тот факт, что значения X схожи между слоями. На основе этого наблюдения мы представляем XQuant-CL, который использует межслойное сходство в X-эмбеддингах для экстремального сжатия. На различных моделях XQuant-CL достигает сокращения памяти до 10 раз по сравнению с базовым уровнем FP16 с ухудшением перплексии всего на 0.01 и до 12.5 раз с ухудшением перплексии на 0.1. XQuant использует быстро растущие вычислительные возможности аппаратных платформ для устранения узкого места в памяти, превосходя современные методы квантования кэша KV и достигая точности, близкой к FP16, на широком спектре моделей.

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.