通过分头卸载降低大模型推理内存占用，实现单卡运行百万级长文本。
HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading
- 按注意力头粒度将键值缓存分批卸载到CPU内存，仅保留部分头在GPU上
- 在100万词元输入下，GPU显存从207GB降至17GB，减少92%
- 无需近似方法即可在24GB显存消费级显卡上运行400万词元推理
基于Transformer的大语言模型在长文本生成中表现优异。随着上下文长度增加，推理时的内存开销主要集中在键值缓存（KV cache）上。本文提出HEADINFER，通过细粒度的分头卸载策略，将KV cache部分移至CPU RAM，避免任何一层完整存储于GPU。该方法仅在GPU上保留部分注意力头的缓存，并动态计算注意力输出。通过屋顶分析证明，HEADINFER在显著降低内存占用的同时保持计算效率。在Llama-3-8B模型、100万词元序列下，KV cache GPU内存从128GB降至1GB，总显存使用从207GB降至17GB，相比BF16基线减少92%。值得注意的是，HEADINFER使8B模型在单张24GB显存消费级显卡（如NVIDIA RTX 4090）上实现400万词元推理，且无需近似方法。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。