HeadInfer: Inferenza Efficiente in Memoria per LLM tramite Offloading a Livello di Testa

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

February 18, 2025

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

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) basati su Transformer dimostrano prestazioni impressionanti nella generazione di contesti lunghi. L'estensione della lunghezza del contesto ha spostato in modo sproporzionato l'impronta di memoria degli LLM durante l'inferenza verso la cache chiave-valore (KV cache). In questo articolo, proponiamo HEADINFER, che scarica la KV cache nella RAM della CPU evitando la necessità di memorizzare completamente la KV cache per qualsiasi livello di transformer sulla GPU. HEADINFER utilizza una strategia di scaricamento granulare e basata sulle teste di attenzione, mantenendo solo la KV cache di specifiche teste di attenzione sulla GPU mentre calcola dinamicamente l'output dell'attenzione. Attraverso un'analisi roofline, dimostriamo che HEADINFER mantiene l'efficienza computazionale riducendo significativamente l'impronta di memoria. Valutiamo HEADINFER sul modello Llama-3-8B con una sequenza di 1 milione di token, riducendo l'impronta di memoria GPU della KV cache da 128 GB a 1 GB e l'utilizzo totale della memoria GPU da 207 GB a 17 GB, ottenendo una riduzione del 92% rispetto all'inferenza di base BF16. In particolare, HEADINFER consente l'inferenza di 4 milioni di token con un modello 8B su una singola GPU consumer con 24 GB di memoria (ad esempio, NVIDIA RTX 4090) senza metodi di approssimazione.

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.