GQA·MQA — KV 캐시를 줄이는 어텐션 헤드 공유

GQA(Grouped-Query Attention)와 MQA(Multi-Query Attention)는 여러 어텐션 헤드가 키(K)·값(V)를 공유하게 만들어 KV 캐시 크기를 줄이는 기법이다. 원래의 멀티헤드 어텐션(MHA)은 헤드마다 자기만의 Q·K·V를 갖는다. MQA는 극단적으로 모든 헤드가 하나의 K·V를 함께 쓰고, GQA는 그 중간에서 몇 개 헤드씩 묶어 K·V를 공유한다. Query는 그대로 여러 개인데 K·V만 줄이는 게 요점이다.

나는 긴 컨텍스트 모델을 서빙하다가 "왜 요즘 모델들은 다 GQA를 쓴다고 적혀 있지?"가 궁금해서 파보기 시작했다. 답은 결국 KV 캐시였다. 생성 속도와 동시 처리량의 발목을 잡는 게 캐시 메모리인데, GQA가 그걸 정직하게 몇 배씩 줄여 준다.

왜 K·V만 줄이나

KV 캐시 크기는 KV 헤드 수에 정비례한다. 헤드가 32개인 MHA라면 32세트의 K·V를 스텝마다 쌓아야 한다. 그런데 어텐션에서 표현력의 상당 부분은 Query가 만든다. Q는 여러 개로 다양하게 두되 K·V만 공유해도 성능 손실이 생각보다 작다는 게 GQA의 발견이다.

MQA는 K·V를 딱 1세트로 통일한다. 캐시가 극적으로 줄지만, 너무 몰아세우면 품질이 눈에 띄게 빠지기도 한다. 그래서 실전에서는 그 사이 어딘가, 예컨대 32개 Q 헤드를 8개 그룹으로 묶어 K·V를 8세트만 두는 식의 GQA가 균형점으로 자리 잡았다.

그림으로 잡는 감

동작을 코드 감각으로 옮기면 이렇다. Q는 헤드별로 다 있고, K·V는 그룹 단위로 만든 뒤 각 Q 헤드가 자기 그룹의 K·V를 참조한다.

# num_heads=32, num_kv_heads=8 이면 그룹당 4개 Q헤드가 K·V 공유
q = proj_q(x)   # (배치, 길이, 32, 헤드차원)
k = proj_k(x)   # (배치, 길이, 8,  헤드차원)  ← 8세트만
v = proj_v(x)   # (배치, 길이, 8,  헤드차원)
# 각 K·V를 그룹 크기(4)만큼 복제해 32개 Q에 맞춘다
k = repeat_kv(k, n=4)   # → 32세트처럼 보이게
v = repeat_kv(v, n=4)
out = attention(q, k, v)   # 이후는 일반 어텐션과 동일

여기서 재밌는 점. 캐시에는 여전히 8세트만 저장한다. 복제는 계산 순간에만 논리적으로 펼치는 것이라, 메모리 이득은 그대로 유지된다. 이 부분을 헷갈리면 "복제하는데 왜 메모리가 줄어?"라는 오해에 빠진다.

어디에 이득이 큰가

이득은 특히 긴 컨텍스트·많은 동시 요청에서 두드러진다. 캐시가 작아지면 같은 VRAM으로 더 긴 문맥을 담거나 더 많은 사용자를 배치에 태울 수 있다. 또 생성 단계는 메모리 대역폭에 묶여 있는데, 읽어야 할 캐시가 줄면 그만큼 토큰 생성이 빨라진다. 요즘 나온 오픈 모델 상당수가 GQA를 기본 채택한 이유다.

MQA는 최대 절약·최대 위험, MHA는 최대 품질·최대 비용. GQA는 그 사이에서 "거의 공짜에 가까운 절약"을 노린 타협이고, 대개 그 타협이 이긴다.

자주 묻는 질문

GQA를 쓰면 모델 품질이 떨어지지 않나요?

그룹 수를 적당히 잡으면 손실이 매우 작다. 이게 GQA가 널리 채택된 핵심 이유다. K·V를 1세트로 몰아넣는 MQA는 품질 저하가 좀 더 보일 수 있지만, 8~4세트 정도의 GQA는 MHA와 체감 차이가 거의 없으면서 캐시를 몇 배 줄인다.

이미 학습된 모델에 GQA를 나중에 적용할 수 있나요?

가능하다. MHA로 학습된 모델의 K·V 헤드를 그룹별로 평균 내는 식으로 초기화한 뒤, 소량의 추가 학습(업트레이닝)으로 GQA 구조로 전환하는 방법이 알려져 있다. 처음부터 GQA로 학습하는 것보다 훨씬 적은 비용으로 캐시 이득을 얻을 수 있다.

플래시 어텐션이나 PagedAttention과 같이 써도 되나요?

서로 보완적이라 함께 쓴다. GQA는 캐시의 "세로 크기(헤드 수)"를 줄이고, PagedAttention은 캐시 메모리의 "배치·할당 낭비"를 줄이며, 플래시 어텐션은 어텐션 연산 자체의 메모리 왕복을 줄인다. 실제 추론 스택은 이 셋을 겹쳐서 쓰는 경우가 많다.

Query 헤드도 같이 줄이면 안 되나요?

Query를 줄이면 어텐션이 볼 수 있는 관점의 다양성이 직접 깎여 표현력 손실이 크다. GQA·MQA가 K·V만 건드리는 이유가 그것이다. 캐시 부담은 K·V에서 오고 표현력은 주로 Q에서 오니, 비용이 큰 쪽만 공유하는 게 영리한 선택이다.