HeadInfer : Inférence efficace en mémoire des LLM par déchargement tête par tête

HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading

February 18, 2025

Auteurs: Cheng Luo, Zefan Cai, Hanshi Sun, Jinqi Xiao, Bo Yuan, Wen Xiao, Junjie Hu, Jiawei Zhao, Beidi Chen, Anima Anandkumar

cs.AI

Résumé

Les grands modèles de langage (LLMs) basés sur l'architecture Transformer démontrent des performances impressionnantes dans la génération de contextes longs. L'extension de la longueur du contexte a déplacé de manière disproportionnée l'empreinte mémoire des LLMs lors de l'inférence vers le cache clé-valeur (KV cache). Dans cet article, nous proposons HEADINFER, qui décharge le KV cache vers la mémoire RAM du CPU tout en évitant la nécessité de stocker entièrement le KV cache pour toute couche Transformer sur le GPU. HEADINFER utilise une stratégie de décharge fine et par tête, ne conservant sur le GPU que le KV cache des têtes d'attention sélectionnées tout en calculant dynamiquement la sortie de l'attention. Grâce à une analyse roofline, nous démontrons que HEADINFER maintient une efficacité computationnelle tout en réduisant significativement l'empreinte mémoire. Nous évaluons HEADINFER sur le modèle Llama-3-8B avec une séquence de 1 million de tokens, réduisant l'empreinte mémoire GPU du KV cache de 128 Go à 1 Go et l'utilisation totale de la mémoire GPU de 207 Go à 17 Go, atteignant ainsi une réduction de 92 % par rapport à l'inférence de référence en BF16. Notamment, HEADINFER permet une inférence sur 4 millions de tokens avec un modèle de 8B sur un seul GPU grand public de 24 Go de mémoire (par exemple, le NVIDIA RTX 4090) sans recourir à des méthodes d'approximation.

English

Transformer-based large language models (LLMs) demonstrate impressive performance in long context generation. Extending the context length has disproportionately shifted the memory footprint of LLMs during inference to the key-value cache (KV cache). In this paper, we propose HEADINFER, which offloads the KV cache to CPU RAM while avoiding the need to fully store the KV cache for any transformer layer on the GPU. HEADINFER employs a fine-grained, head-wise offloading strategy, maintaining only selective attention heads KV cache on the GPU while computing attention output dynamically. Through roofline analysis, we demonstrate that HEADINFER maintains computational efficiency while significantly reducing memory footprint. We evaluate HEADINFER on the Llama-3-8B model with a 1-million-token sequence, reducing the GPU memory footprint of the KV cache from 128 GB to 1 GB and the total GPU memory usage from 207 GB to 17 GB, achieving a 92% reduction compared to BF16 baseline inference. Notably, HEADINFER enables 4-million-token inference with an 8B model on a single consumer GPU with 24GB memory (e.g., NVIDIA RTX 4090) without approximation methods.