HeadInfer: Inferencia Eficiente en Memoria para Modelos de Lenguaje mediante Descarga por Cabezales

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

Resumen

Los modelos de lenguaje de gran escala (LLMs) basados en Transformers demuestran un rendimiento impresionante en la generación de contextos largos. La extensión de la longitud del contexto ha desplazado de manera desproporcionada la huella de memoria de los LLMs durante la inferencia hacia la caché de clave-valor (KV cache). En este artículo, proponemos HEADINFER, que descarga la KV cache a la memoria RAM de la CPU mientras evita la necesidad de almacenar completamente la KV cache para cualquier capa del Transformer en la GPU. HEADINFER emplea una estrategia de descarga granular a nivel de cabezas, manteniendo solo la KV cache de cabezas de atención selectivas en la GPU mientras calcula dinámicamente la salida de atención. Mediante un análisis de techo de rendimiento, demostramos que HEADINFER mantiene la eficiencia computacional mientras reduce significativamente la huella de memoria. Evaluamos HEADINFER en el modelo Llama-3-8B con una secuencia de 1 millón de tokens, reduciendo la huella de memoria de la GPU de la KV cache de 128 GB a 1 GB y el uso total de memoria de la GPU de 207 GB a 17 GB, logrando una reducción del 92% en comparación con la inferencia de referencia BF16. Notablemente, HEADINFER permite la inferencia de 4 millones de tokens con un modelo de 8B en una sola GPU de consumo con 24 GB de memoria (por ejemplo, NVIDIA RTX 4090) sin métodos de aproximación.

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.