HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading
Fuente:
arXiv
Guardado en:
| Autores principales: | , , , , , , , , , |
|---|---|
| Formato: | Preprint |
| Publicado: |
2025
|
| Materias: | |
| Acceso en línea: | |
| Etiquetas: |
Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
|
| _version_ | 1866929718922575872 |
|---|---|
| author | Luo, Cheng Cai, Zefan Sun, Hanshi Xiao, Jinqi Yuan, Bo Xiao, Wen Hu, Junjie Zhao, Jiawei Chen, Beidi Anandkumar, Anima |
| author_facet | Luo, Cheng Cai, Zefan Sun, Hanshi Xiao, Jinqi Yuan, Bo Xiao, Wen Hu, Junjie Zhao, Jiawei Chen, Beidi Anandkumar, Anima |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_12574 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading Luo, Cheng Cai, Zefan Sun, Hanshi Xiao, Jinqi Yuan, Bo Xiao, Wen Hu, Junjie Zhao, Jiawei Chen, Beidi Anandkumar, Anima Machine Learning Artificial Intelligence 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. |
| title | HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2502.12574 |