HeadInfer: Geheugenefficiënte LLM-inferentie door head-wise offloading

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

Samenvatting

Transformer-gebaseerde grote taalmodellen (LLMs) tonen indrukwekkende prestaties in het genereren van lange contexten. Het uitbreiden van de contextlengte heeft onevenredig de geheugenvoetafdruk van LLMs tijdens inferentie verschoven naar de key-value cache (KV-cache). In dit artikel stellen we HEADINFER voor, dat de KV-cache uitbesteedt aan het CPU-RAM terwijl het vermijdt om de KV-cache volledig op te slaan voor enige transformer-laag op de GPU. HEADINFER maakt gebruik van een fijnmazige, head-wise uitbestedingsstrategie, waarbij alleen selectieve aandacht heads KV-cache op de GPU wordt behouden terwijl de aandachtoutput dynamisch wordt berekend. Door middel van roofline-analyse tonen we aan dat HEADINFER de rekenkundige efficiëntie behoudt terwijl het de geheugenvoetafdruk aanzienlijk vermindert. We evalueren HEADINFER op het Llama-3-8B-model met een sequentie van 1 miljoen tokens, waarbij de GPU-geheugenvoetafdruk van de KV-cache wordt teruggebracht van 128 GB naar 1 GB en het totale GPU-geheugengebruik van 207 GB naar 17 GB, wat een reductie van 92% oplevert in vergelijking met BF16 baseline-inferentie. Opmerkelijk is dat HEADINFER inferentie met 4 miljoen tokens mogelijk maakt met een 8B-model op een enkele consumenten-GPU met 24GB geheugen (bijv. NVIDIA RTX 4090) zonder benaderingsmethoden.

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.