HeadInfer: ヘッド単位オフロードによるメモリ効率の良い大規模言語モデル推論

HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading

February 18, 2025

著者: Cheng Luo, Zefan Cai, Hanshi Sun, Jinqi Xiao, Bo Yuan, Wen Xiao, Junjie Hu, Jiawei Zhao, Beidi Chen, Anima Anandkumar

cs.AI

要旨

Transformerベースの大規模言語モデル（LLM）は、長文脈生成において驚異的な性能を発揮します。文脈長の拡大に伴い、推論時のLLMのメモリ使用量は、キー・バリューキャッシュ（KVキャッシュ）に偏って増加しています。本論文では、HEADINFERを提案します。これは、KVキャッシュをCPU RAMにオフロードしつつ、GPU上でどのTransformer層のKVキャッシュも完全に保持する必要を回避するものです。HEADINFERは、細粒度のヘッド単位のオフロード戦略を採用し、GPU上では選択的なアテンションヘッドのKVキャッシュのみを保持し、アテンション出力を動的に計算します。ルーフライン分析を通じて、HEADINFERが計算効率を維持しつつ、メモリ使用量を大幅に削減することを示します。Llama-3-8Bモデルを用いて100万トークンのシーケンスで評価を行った結果、KVキャッシュのGPUメモリ使用量を128GBから1GBに、総GPUメモリ使用量を207GBから17GBに削減し、BF16ベースライン推論と比較して92%の削減を達成しました。特に、HEADINFERは、近似手法を用いることなく、24GBメモリを搭載した単一のコンシューマーGPU（例：NVIDIA RTX 4090）上で8Bモデルの400万トークン推論を可能にします。

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.