이 논문은 대규모 언어 모델(LLM) 추론 시 발생하는 KV 캐시 문제를 해결하기 위한 SQUEEZEATTENTION 알고리즘을 제안한다.
먼저 저자들은 각 주의 층의 중요도를 코사인 유사도로 측정하여 분석했다. 그 결과 주의 층마다 중요도가 다르다는 것을 발견했다. 이를 바탕으로 SQUEEZEATTENTION은 레이어 차원에서 KV 캐시 예산을 동적으로 재할당하여 최적화한다.
구체적으로 SQUEEZEATTENTION은 먼저 입력 프롬프트 처리 중 각 레이어의 코사인 유사도를 측정하여 레이어들을 중요도 그룹으로 클러스터링한다. 그 후 중요도가 낮은 그룹의 레이어에 더 적은 캐시 예산을 할당하고, 중요도가 높은 그룹의 레이어에 더 많은 예산을 할당한다.
이렇게 레이어 차원에서 최적화된 KV 캐시 예산을 바탕으로, SQUEEZEATTENTION은 기존의 토큰 기반 압축 알고리즘(Sliding Window, H2O 등)을 각 레이어에 적용하여 최종적인 KV 캐시를 압축한다.
실험 결과, SQUEEZEATTENTION은 다양한 LLM 모델과 데이터셋에서 기존 알고리즘 대비 30-70%의 메모리 절감과 최대 2.2배의 처리량 향상을 달성했다.
To Another Language
from source content
arxiv.org
Key Insights Distilled From
by Zihao Wang,S... at arxiv.org 04-09-2024
https://arxiv.org/pdf/2404.04793.pdfDeeper Inquiries