HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading
Authors: name: Cheng Luo, name: Zefan Cai, name: Hanshi Sun, name: Jinqi Xiao, name: Bo Yuan, name: Wen Xiao, name: Junjie Hu, name: Jiawei Zhao, name: Beidi Chen, name: Anima Anandkumar

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.