멀티쿼리·그룹쿼리 어텐션(MQA·GQA): LLM 추론 속도를 높이는 KV 캐시 절약법

핵심 메시지

• KV 캐시는 토큰마다의 어텐션 계산을 줄여주지만, 그 대가로 메모리와 메모리 대역폭이 새로운 병목이 된다.
• MQA(멀티쿼리 어텐션)는 모든 어텐션 헤드가 키와 값을 하나로 공유해 KV 캐시 크기를 크게 줄인다.
• 키·값을 공유하면 표현력이 떨어져 성능 손실이 생기므로, 그 절충안으로 GQA가 등장했다.
• GQA(그룹쿼리 어텐션)는 쿼리 헤드를 그룹으로 나눠 그룹 안에서만 키·값을 공유함으로써 속도와 성능을 모두 챙긴다.
• 구현상의 변경은 작다 — KV 투영 차원과 KV 캐시 모양을 바꾸고, 공유된 키·값을 어텐션 전에 펼쳐주기만 하면 된다.

쉽게 이해하기

자기회귀 생성에서 KV 캐시는 단계마다의 어텐션 계산을 O(n²)에서 O(n)으로 줄여준다. 그러나 이 이점에는 대가가 따른다. 긴 문맥, 큰 배치, 깊은 모델일수록 캐시 자체가 매우 커지고, 이는 모델 가중치와 활성값이 이미 차지한 GPU 메모리를 함께 갉아먹어 추론 효율을 떨어뜨린다.

문제는 메모리 용량만이 아니다. 최신 GPU는 계산은 매우 빠르지만, 추론 중에는 종종 계산이 아니라 메모리 대역폭에 발목이 잡힌다. 토큰을 하나 생성할 때마다 모든 층에서 큰 KV 캐시 데이터를 VRAM에서 빠른 온칩 메모리로 거듭 옮겨야 하는데, 정작 어텐션 계산보다 이 데이터 이동에 더 많은 시간이 든다. 그래서 KV 캐시를 줄이는 일이 그토록 중요하다.

표준 멀티헤드 어텐션(MHA)에서는 헤드마다 독립적인 Q·K·V 투영을 가진다. 반면 MQA에서는 쿼리 투영은 그대로 두되 모든 헤드가 키와 값을 단 하나로 공유한다. 헤드별로 키·값을 따로 저장하지 않으니 캐시가 급격히 줄지만, 모든 관계를 하나의 공유 표현 공간에 욱여넣어야 하므로 표현력은 떨어진다.

GQA는 그 중간 지대를 노린다. 모든 헤드가 키·값을 공유(MQA)하지도, 헤드마다 따로 갖지도(MHA) 않고, 쿼리 헤드를 여러 그룹으로 나눠 같은 그룹끼리만 키·값을 공유한다. 그룹 수를 1로 하면 MQA, 헤드 수만큼 늘리면 MHA가 되므로, GQA는 두 방식을 잇는 연속적인 스펙트럼 위에 놓인다.

효과는 수치로 분명하다. 64개 헤드에 배치 크기 4인 최신 LLM에서 긴 문맥의 KV 캐시는 약 80GB까지 불어날 수 있는데, MQA로 바꾸면 약 1.25GB 수준으로 줄어든다. GQA 논문은 MQA가 속도를 크게 높이지만 성능 저하를 동반하는 반면, GQA는 속도 이점 대부분을 유지하면서 잃었던 성능을 거의 회복한다고 보고했다. 또한 처음부터 MQA로 학습하면 불안정하기 때문에, 기존 MHA 체크포인트의 헤드들을 평균 풀링으로 합쳐 MQA·GQA로 변환한 뒤 소량만 추가 학습하는 방법이 제안됐다.

주요 인사이트

• 추론의 진짜 병목은 종종 연산이 아니라 메모리 대역폭이다. 계산 자체보다 KV 캐시를 옮기는 데 시간이 더 들어, 강력한 연산 코어가 데이터를 기다리며 놀게 된다.
• 그룹 수를 1로 두면 MQA, 헤드 수만큼 두면 MHA가 된다. GQA는 둘 사이를 잇는 연속적인 절충점으로 이해하면 쉽다.
• 처음부터 MQA로 학습하면 특히 긴 문맥에서 불안정해진다. 그래서 이미 학습된 MHA 모델의 키·값을 평균 풀링으로 합쳐 변환한 뒤 짧게 추가 학습하는 방식이 더 안정적이다.
• GQA가 MQA보다 널리 채택되는 이유는 분명하다. 성능을 대부분 회복하면서도 추론 속도 이점의 상당 부분을 그대로 유지하기 때문이다.

자주 묻는 질문

관련 AI 소식