HeadInfer: Speichereffiziente LLM-Inferenz durch kopfweise Auslagerung

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

February 18, 2025

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

cs.AI

Zusammenfassung

Transformer-basierte große Sprachmodelle (LLMs) zeigen beeindruckende Leistungen bei der Generierung langer Kontexte. Die Erweiterung der Kontextlänge hat den Speicherbedarf von LLMs während der Inferenz unverhältnismäßig stark auf den Key-Value-Cache (KV-Cache) verlagert. In diesem Artikel präsentieren wir HEADINFER, das den KV-Cache in den CPU-RAM auslagert, während es vermeidet, den KV-Cache für jede Transformer-Schicht vollständig auf der GPU zu speichern. HEADINFER verwendet eine feingranulare, kopfweise Auslagerungsstrategie, bei der nur ausgewählte Attention-Heads des KV-Caches auf der GPU gehalten werden, während die Attention-Ausgabe dynamisch berechnet wird. Durch eine Roofline-Analyse zeigen wir, dass HEADINFER die Recheneffizienz beibehält und gleichzeitig den Speicherbedarf erheblich reduziert. Wir evaluieren HEADINFER am Llama-3-8B-Modell mit einer 1-Million-Token-Sequenz und reduzieren den GPU-Speicherbedarf des KV-Caches von 128 GB auf 1 GB sowie den gesamten GPU-Speicherverbrauch von 207 GB auf 17 GB, was einer Reduktion von 92 % im Vergleich zur BF16-Baseline-Inferenz entspricht. Bemerkenswerterweise ermöglicht HEADINFER die Inferenz von 4 Millionen Tokens mit einem 8B-Modell auf einer einzelnen Consumer-GPU mit 24 GB Speicher (z. B. NVIDIA RTX 4090) ohne Näherungsmethoden.

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.