HeadInfer: Inferência Eficiente em Memória para LLMs por meio de Descarregamento por Cabeças

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

February 18, 2025

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

cs.AI

Resumo

Modelos de linguagem de grande escala (LLMs) baseados em Transformers demonstram desempenho impressionante na geração de contextos longos. A extensão do comprimento do contexto deslocou desproporcionalmente a pegada de memória dos LLMs durante a inferência para o cache de chave-valor (KV cache). Neste artigo, propomos o HEADINFER, que transfere o KV cache para a memória RAM da CPU, evitando a necessidade de armazenar completamente o KV cache para qualquer camada do Transformer na GPU. O HEADINFER emprega uma estratégia de transferência granular e orientada por cabeças, mantendo apenas o cache KV de cabeças de atenção seletivas na GPU enquanto calcula dinamicamente a saída de atenção. Através da análise de roofline, demonstramos que o HEADINFER mantém a eficiência computacional enquanto reduz significativamente a pegada de memória. Avaliamos o HEADINFER no modelo Llama-3-8B com uma sequência de 1 milhão de tokens, reduzindo a pegada de memória da GPU do KV cache de 128 GB para 1 GB e o uso total de memória da GPU de 207 GB para 17 GB, alcançando uma redução de 92% em comparação com a inferência de linha de base BF16. Notavelmente, o HEADINFER permite a inferência de 4 milhões de tokens com um modelo de 8B em uma única GPU de consumo com 24GB de memória (por exemplo, NVIDIA RTX 4090) sem métodos de aproximação.

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.