title: HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading
summary: 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.
paper_authors._id: 67b547f555d0424a31b9c384
paper_authors.name: Cheng Luo
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-02-19T09:40:25.130Z
paper_authors._id: 64cb48f7667f4f808535107e
paper_authors.avatarUrl: /avatars/8f77f378ad665b246e1ea3aaba2153ae.svg
paper_authors.fullname: chengluo
paper_authors.type: user
paper_authors.user: wdlctc
paper_authors._id: 67b547f555d0424a31b9c385
paper_authors.name: Zefan Cai
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-02-19T09:40:47.077Z
paper_authors._id: 64b15284372d4340772a3dca
paper_authors.avatarUrl: /avatars/417d5f1bc1bcb5e4d5de6169673c2cf7.svg
paper_authors.fullname: Zefan Cai
paper_authors.type: user
paper_authors.user: ZefanCai
paper_authors._id: 67b547f555d0424a31b9c386
paper_authors.name: Hanshi Sun
paper_authors._id: 67b547f555d0424a31b9c387
paper_authors.name: Jinqi Xiao
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-02-19T09:41:01.931Z
paper_authors._id: 64c15c5bea792b1950e302e4
paper_authors.avatarUrl: /avatars/51f84365cc08a1dcd5da70968389aed2.svg
paper_authors.fullname: Jinqi Xiao
paper_authors.type: user
paper_authors.user: jinqixiao
paper_authors._id: 67b547f555d0424a31b9c388
paper_authors.name: Bo Yuan
paper_authors._id: 67b547f555d0424a31b9c389
paper_authors.name: Wen Xiao
paper_authors._id: 67b547f555d0424a31b9c38a
paper_authors.name: Junjie Hu
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-02-19T09:41:18.304Z
paper_authors._id: 675f8271a63fff7b5bcbc478
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/no-auth/9tJn7NyzLMreCJVH4wRho.png
paper_authors.fullname: Junjie Hu
paper_authors.type: user
paper_authors.user: junjiehu
paper_authors._id: 67b547f555d0424a31b9c38b
paper_authors.name: Jiawei Zhao
paper_authors._id: 67b547f555d0424a31b9c38c
paper_authors.name: Beidi Chen
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-02-19T09:39:20.563Z
paper_authors._id: 64b732f832403871593e082c
paper_authors.avatarUrl: /avatars/dd21932b0c167131ee7545a622c46c3c.svg
paper_authors.fullname: Beidi Chen
paper_authors.type: user
paper_authors.user: beidic
paper_authors._id: 67b547f555d0424a31b9c38d
paper_authors.name: Anima Anandkumar
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-02-19T09:39:15.091Z
paper_authors._id: 6532920b3e385cfc6002938d
paper_authors.avatarUrl: /avatars/cb9cc6d2733031582c83f56dc6cd1dd5.svg
paper_authors.fullname: Anima Anandkumar
paper_authors.type: user
paper_authors.user: animakumar
paper_title: HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading
paper_summary: 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.
paper_ai_keywords: Transformer-based large language models (LLMs) key-value cache (KV cache) HEADINFER fine-grained, head-wise offloading strategy roofline analysis Llama-3-8B model 1-million-token sequence 128 GB GPU memory footprint reduction 207 GB to 17 GB GPU memory usage reduction 92% reduction compared to BF16 baseline inference 4-million-token inference