FairKV: Balancing Per-Head KV Cache for Fast Multi-GPU Inference
Created by
Haebom
저자
Bingzhe Zhao, Ke Cheng, Aomufei Yuan, Yuxuan Tian, Ruiguang Zhong, Chengchen Hu, Tong Yang, Lian Yu
개요
본 논문은 Transformer 모델에서 KV 캐시 압축의 불균형적인 메모리 할당으로 인한 다중 GPU 추론 시의 부하 불균형 문제를 해결하기 위해 FairKV를 제안한다. 기존 최첨단 KV 캐시 압축 방법들은 각 어텐션 헤드에 대한 메모리 할당을 동적으로 조정하는 불균형적인 알고리즘을 사용하여 단일 GPU 환경에서 우수한 성능을 달성하지만, 다중 GPU 환경에서는 특정 GPU에 과부하가 걸리는 문제가 발생한다. FairKV는 Fair-Copying 기법을 통해 메모리 사용량이 많은 어텐션 헤드의 작은 부분집합을 데이터 병렬 처리 방식으로 GPU에 복제하여 부하 불균형을 완화한다. LLaMA 70b 및 Mistral 24b 모델 실험 결과, FairKV는 표준 텐서 병렬 추론에 비해 처리량을 1.66배 향상시키는 것으로 나타났다.
시사점, 한계점
•
시사점:
◦
다중 GPU 환경에서의 Transformer 모델 추론 성능 향상을 위한 효과적인 방법 제시.
◦
불균형적인 KV 캐시 압축으로 인한 부하 불균형 문제 해결.
◦
데이터 병렬 처리를 활용한 Fair-Copying 기법의 효용성 증명.
◦
LLaMA 70b 및 Mistral 24b 모델에서의 성능 향상을 통해 실제 적용 가능성 확인.
•
한계점:
◦
Fair-Copying 기법의 추가적인 메모리 오버헤드 발생 가능성.
◦
다양한 모델 및 하드웨어 환경에 대한 추가적인 실험 필요.
◦
Fair-Copying 기법의 최적 매개변수(복제할 어텐션 헤드의 크기 등) 결정에 대한 추가적인 연구 필요.
