HeadInfer: 헤드 단위 오프로딩을 통한 메모리 효율적인 LLM 추론

초록

Transformer 기반의 대규모 언어 모델(LLM)은 긴 문맥 생성에서 인상적인 성능을 보여줍니다. 문맥 길이를 확장함에 따라, 추론 과정에서 LLM의 메모리 사용량은 키-값 캐시(KV 캐시)에 불균형적으로 집중되었습니다. 본 논문에서는 HEADINFER를 제안하며, 이는 KV 캐시를 CPU RAM으로 오프로드하면서도 GPU 상의 어떤 Transformer 레이어에서도 KV 캐시를 완전히 저장할 필요를 없앱니다. HEADINFER는 세분화된 헤드 단위 오프로딩 전략을 사용하여, GPU 상에서 선택적인 어텐션 헤드의 KV 캐시만 유지하면서 어텐션 출력을 동적으로 계산합니다. Roofline 분석을 통해 HEADINFER가 계산 효율성을 유지하면서도 메모리 사용량을 크게 줄인다는 것을 입증합니다. HEADINFER를 100만 토큰 시퀀스로 Llama-3-8B 모델에서 평가한 결과, 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.

HeadInfer: 헤드 단위 오프로딩을 통한 메모리 효율적인 LLM 추론

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

초록

Summary

Support