그룹 쿼리 어텐션

Grouped-query attention, GQA; 그룹 쿼리 어텐션
쿼리 헤드를 몇 개의 그룹으로 묶고 그룹마다 키·값 헤드 하나를 공유하게 해, KV 캐시와 디코딩 메모리 대역폭을 줄이는 멀티 헤드 어텐션 변형

그룹 쿼리 어텐션(Grouped-query attention, GQA)은 멀티 헤드 어텐션(MHA)의 쿼리 헤드 $H$ 개를 $G$ 개 그룹으로 나누고, 같은 그룹의 쿼리 헤드들이 키·값 헤드 한 쌍을 함께 쓰게 하는 어텐션 구조다. 2023년 구글 리서치의 에인슬리(Joshua Ainslie) 등이 제안했고 EMNLP 2023에 실렸다.^[1] 키·값 헤드를 하나만 두는 멀티 쿼리 어텐션(Multi-query attention, MQA)과 헤드마다 따로 두는 MHA의 중간이다. GQA-1은 MQA와 같고, GQA- $H$ 는 MHA와 같다.

목적은 학습이 아니라 추론이다. 자기회귀 디코딩은 토큰 하나를 만들 때마다 지금까지의 키·값 전체(KV 캐시)를 메모리에서 읽어야 하므로, 계산보다 메모리 대역폭이 속도를 정한다. 키·값 헤드를 $H / G$ 배 줄이면 캐시 크기와 읽을 데이터가 그만큼 준다. 메타 라마 2의 34B·70B, 라마 3 전 모델, 미스트랄 7B, 구글 젬마 2 등 2023년 이후 공개 모델 상당수가 GQA를 썼다.^[2]^[3]^[4]^[5]

배경: 디코딩은 메모리에 묶여 있다

학습이나 프롬프트 처리(프리필)는 시퀀스 전체를 한 번에 행렬 곱으로 처리하므로 GPU 연산 장치가 바쁘다. 반면 생성 단계는 쿼리가 토큰 하나뿐이라, 큰 키·값 텐서를 읽어 와서 작은 계산만 하고 끝난다. 섀지어(Noam Shazeer)는 2019년 논문에서 증분 디코딩이 느린 원인을 "큰 키·값 텐서를 반복해서 읽는 메모리 대역폭 비용"으로 짚고, 이를 줄이려고 MQA를 제안했다.^[6]

멀티 쿼리 어텐션 (MQA)

MQA는 쿼리 투영은 헤드마다 따로 두고, 키·값 투영은 하나만 둔다. 모든 쿼리 헤드가 같은 키·값을 본다. WMT14 영어-독일어 번역 실험(TPUv2, 시퀀스 길이 128)에서 디코더의 토큰당 추론 시간이 46μs에서 3.8μs로 줄었고, BLEU는 기준 모델과 비슷했다(빔 4 테스트 28.4 대 28.5). 매개변수 수를 맞추려고 FFN 은닉 크기를 4096에서 5440으로 늘린 조건이다.^[6]

구글 PaLM(2022)이 MQA를 썼고, 젬마 1의 2B 모델도 소규모에서는 MQA가 잘 동작한다는 절제 실험을 근거로 MQA를 썼다(7B는 MHA).^[7]^[8] 다만 GQA 논문은 MQA가 품질 저하와 학습 불안정을 일으킬 수 있다고 지적했다.^[1]

구조

쿼리 헤드 $H$ 개, 키·값 헤드 $G$ 개( $H$ 가 $G$ 로 나누어떨어져야 함)일 때, 쿼리 헤드 $i$ 는 키·값 헤드 $g (i) = ⌊ i \cdot G / H ⌋$ 를 쓴다.

${h e a d}_{i} = softmax (\frac{(X W_{i}^{Q}) (X W_{g (i)}^{K})^{⊤}}{\sqrt{d_{h}}}) X W_{g (i)}^{V}$

방식	키·값 헤드 수	토큰·층당 KV 캐시(원소 수)	키·값 투영 매개변수
MHA	$H$	$2 H d_{h}$	$2 D \cdot H d_{h}$
GQA- $G$	$G$	$2 G d_{h}$	$2 D \cdot G d_{h}$
MQA	1	$2 d_{h}$	$2 D \cdot d_{h}$

GQA 논문은 그룹 수를 중간으로 두는 이유를 이렇게 설명한다. 모델이 커지면 헤드 수도 늘어나므로 MQA는 큰 모델일수록 더 과격하게 깎는 셈이지만, GQA는 모델 크기와 상관없이 같은 비율로 줄일 수 있다. 또 큰 모델은 KV 캐시가 모델 차원에 비례하는 반면 연산량과 매개변수는 그 제곱에 비례하므로 어텐션의 대역폭 부담이 상대적으로 작다. 그리고 텐서 병렬로 여러 장치에 나눌 때 MQA는 하나뿐인 키·값 헤드를 장치마다 복제해야 하지만, GQA는 그룹을 장치에 나눠 배치할 수 있어 낭비가 없다.^[1]

메모리 계산 예

라마 3 70B는 층 80개, 모델 차원 8192, 쿼리 헤드 64개, 키·값 헤드 8개(헤드 차원 128)다.^[3] BF16(2바이트)로 128K 토큰 한 시퀀스의 KV 캐시를 계산하면 다음과 같다.

가정	계산	크기
MHA였다면(키·값 헤드 64)	2 × 80 × 64 × 128 × 131072 × 2바이트	320 GiB
실제 GQA-8	2 × 80 × 8 × 128 × 131072 × 2바이트	40 GiB
MQA였다면	2 × 80 × 1 × 128 × 131072 × 2바이트	5 GiB

라마 3 보고서는 GQA를 "추론 속도를 높이고 디코딩 중 키-값 캐시 크기를 줄이려고" 썼다고 밝혔다.^[3] 미스트랄 7B는 쿼리 헤드 32개에 키·값 헤드 8개,^[4] 젬마 2는 2B·9B·27B 모두 쿼리 헤드 2개가 키·값 헤드 하나를 공유한다.^[5]

업트레이닝: MHA 체크포인트를 GQA로

GQA 논문의 다른 기여는 이미 MHA로 학습된 모델을 처음부터 다시 학습하지 않고 바꾸는 방법이다.^[1]

체크포인트 변환: 그룹에 속한 원래 키 헤드들의 투영 행렬을 평균(mean pooling)해 그룹의 키 헤드로 삼는다. 값도 같다. 첫 번째 헤드만 고르거나 새로 무작위 초기화하는 것보다 평균이 나았다.
추가 사전 학습: 원래 사전 학습 단계의 일부 비율 $α$ 만큼 같은 설정으로 더 학습한다. $α = 0.05$ (원래 계산량의 5%)일 때 T5-XXL 기준 약 600 TPUv3 칩-일이 들었다.

T5.1.1 XXL을 5% 업트레이닝한 결과(요약·번역·질의응답 7개 과제 평균, TPUv4 칩당 표본 하나 추론 시간)는 다음과 같다. 인코더 셀프 어텐션에는 적용하지 않고 디코더 셀프 어텐션과 교차 어텐션에만 적용했다.^[1]

모델	추론 시간(초)	평균 점수
MHA-Large	0.37	46.0
MHA-XXL	1.51	47.2
MQA-XXL (업트레이닝)	0.24	46.6
GQA-8-XXL (업트레이닝)	0.28	47.1

GQA-8은 MHA-XXL과 거의 같은 품질을 MQA에 가까운 속도로 냈다. 변환 코드는 간단하다.

import torch

def mha_to_gqa_kv(w_k: torch.Tensor, n_heads: int, n_kv_heads: int) -> torch.Tensor:
    """w_k: (n_heads*head_dim, d_model) 키(또는 값) 투영 가중치 -> (n_kv_heads*head_dim, d_model)"""
    head_dim = w_k.shape[0] // n_heads
    group = n_heads // n_kv_heads
    w = w_k.view(n_kv_heads, group, head_dim, -1)   # 연속한 헤드끼리 한 그룹
    return w.mean(dim=1).reshape(n_kv_heads * head_dim, -1)

구현

학습·프리필 때는 키·값 헤드를 그룹 크기만큼 복제해 일반 MHA처럼 계산하면 된다. 복제는 expand로 하면 메모리를 새로 잡지 않는다. 허깅 페이스 Transformers의 repeat_kv, 메타 공식 코드의 같은 이름 함수가 이 일을 한다.^[9]^[10] 허깅 페이스 설정에서는 num_key_value_heads가 키·값 헤드 수다. PyTorch 2.5부터는 scaled_dot_product_attention에 enable_gqa=True를 주면 복제 없이 바로 계산한다.^[11]

import torch
import torch.nn as nn
import torch.nn.functional as F

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__()
        assert n_heads % n_kv_heads == 0
        self.h, self.h_kv = n_heads, n_kv_heads
        self.d = d_model // n_heads
        self.wq = nn.Linear(d_model, n_heads * self.d, bias=False)
        self.wk = nn.Linear(d_model, n_kv_heads * self.d, bias=False)   # 키·값 투영이 작아진다
        self.wv = nn.Linear(d_model, n_kv_heads * self.d, bias=False)
        self.wo = nn.Linear(n_heads * self.d, d_model, bias=False)

    def forward(self, x):
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.h, self.d).transpose(1, 2)      # (B, H, T, d)
        k = self.wk(x).view(B, T, self.h_kv, self.d).transpose(1, 2)   # (B, G, T, d): KV 캐시에 들어가는 것
        v = self.wv(x).view(B, T, self.h_kv, self.d).transpose(1, 2)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)  # PyTorch 2.5+
        # 2.5 미만: k = k.repeat_interleave(self.h // self.h_kv, dim=1) 후 enable_gqa 없이 호출
        return self.wo(y.transpose(1, 2).reshape(B, T, -1))

복제 순서에 주의한다. 위 코드와 라마 구현은 연속한 쿼리 헤드들이 같은 키·값 헤드를 쓰는 배치(repeat_interleave)다. repeat로 복제하면 헤드 대응이 달라져 기존 체크포인트와 맞지 않는다.

GQA를 직접 지원하는 어텐션 커널은 키·값을 복제하지 않고 한 키·값 헤드를 읽어 그룹의 쿼리 헤드 여러 개에 한꺼번에 쓰므로, 읽은 바이트당 연산량(산술 강도)이 그룹 크기만큼 올라간다. 이것이 GQA가 캐시 크기 이상으로 속도에 기여하는 이유다.

한계와 반론

품질 손실: 키·값 표현력이 줄어든다. 딥시크는 같은 조건(1.33조 토큰, 약 70억 매개변수)으로 학습한 밀집 모델에서 MMLU가 MHA 45.2, GQA-8 41.2, MQA 37.9로 MHA가 뚜렷하게 앞섰다고 보고했다. 이 결과를 근거로 키·값을 헤드 수 대신 저차원 압축으로 줄이는 멀티 헤드 잠재 어텐션(MLA)을 만들었다.^[12]
텐서 병렬 제약: 키·값 헤드 수가 텐서 병렬 장치 수보다 적으면 헤드를 나눌 수 없어 결국 복제해야 한다. 키·값 헤드 8개 모델을 16개 GPU로 나누면 각 헤드가 두 장치에 중복된다.
긴 문맥에서는 부족: 캐시가 MHA의 $G / H$ 로 줄 뿐 여전히 길이에 비례해 커진다. 수십만~수백만 토큰을 다루는 모델은 GQA에 KV 캐시 가지치기, 캐시 양자화, 슬라이딩 윈도우, 선형 어텐션 층 혼합 등을 겹친다.

실무 선택 기준

새로 사전 학습하는 디코더 모델이면 GQA가 무난한 기본값이다. 앞서 본 공개 모델들은 쿼리 헤드 2~8개당 키·값 헤드 하나를 두었다(라마 3: 8개당 1개, 미스트랄 7B: 4개당 1개, 젬마 2: 2개당 1개).
이미 MHA로 학습한 모델의 추론 비용이 문제라면 평균 변환 후 원 학습량의 수 % 업트레이닝을 고려한다.
메모리가 아주 빠듯한 소형 모델이나 온디바이스 모델은 MQA도 선택지다.
키·값 헤드 수는 배포할 텐서 병렬 크기로 나누어떨어지게 정한다.

같이 보기

각주

↑ ^1.0 ^1.1 ^1.2 ^1.3 ^1.4 Joshua Ainslie 외, 〈GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints〉, EMNLP 2023. arXiv:2305.13245
↑ Hugo Touvron 외, 〈Llama 2: Open Foundation and Fine-Tuned Chat Models〉, 2023. arXiv:2307.09288
↑ ^3.0 ^3.1 ^3.2 Llama Team, AI @ Meta, 〈The Llama 3 Herd of Models〉, 2024. arXiv:2407.21783
↑ ^4.0 ^4.1 Albert Q. Jiang 외, 〈Mistral 7B〉, 2023. arXiv:2310.06825
↑ ^5.0 ^5.1 Gemma Team, 〈Gemma 2: Improving Open Language Models at a Practical Size〉, 2024. arXiv:2408.00118
↑ ^6.0 ^6.1 Noam Shazeer, 〈Fast Transformer Decoding: One Write-Head is All You Need〉, 2019. arXiv:1911.02150
↑ Aakanksha Chowdhery 외, 〈PaLM: Scaling Language Modeling with Pathways〉, 2022. arXiv:2204.02311
↑ Gemma Team, 〈Gemma: Open Models Based on Gemini Research and Technology〉, 2024. arXiv:2403.08295
↑ huggingface/transformers, modeling_llama.py
↑ meta-llama/llama-models, models/llama3/model.py
↑ PyTorch 2.5 문서, torch.nn.functional.scaled_dot_product_attention
↑ DeepSeek-AI, 〈DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model〉, 2024. arXiv:2405.04434

[gqa-1] 1.0 ^1.1 ^1.2 ^1.3 ^1.4 Joshua Ainslie 외, 〈GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints〉, EMNLP 2023. arXiv:2305.13245

[llama2-2] Hugo Touvron 외, 〈Llama 2: Open Foundation and Fine-Tuned Chat Models〉, 2023. arXiv:2307.09288

[llama3-3] 3.0 ^3.1 ^3.2 Llama Team, AI @ Meta, 〈The Llama 3 Herd of Models〉, 2024. arXiv:2407.21783

[mistral-4] 4.0 ^4.1 Albert Q. Jiang 외, 〈Mistral 7B〉, 2023. arXiv:2310.06825

[gemma2-5] 5.0 ^5.1 Gemma Team, 〈Gemma 2: Improving Open Language Models at a Practical Size〉, 2024. arXiv:2408.00118

[mqa-6] 6.0 ^6.1 Noam Shazeer, 〈Fast Transformer Decoding: One Write-Head is All You Need〉, 2019. arXiv:1911.02150

[palm-7] Aakanksha Chowdhery 외, 〈PaLM: Scaling Language Modeling with Pathways〉, 2022. arXiv:2204.02311

[gemma1-8] Gemma Team, 〈Gemma: Open Models Based on Gemini Research and Technology〉, 2024. arXiv:2403.08295

[hfllama-9] uggingface/transformers, modeling_llama.py

[metallama-10] ta-llama/llama-models, models/llama3/model.py

[sdpa-11] PyTorch 2.5 문서, torch.nn.functional.scaled_dot_product_attention

[dsv2-12] DeepSeek-AI, 〈DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model〉, 2024. arXiv:2405.04434

[1]

[2]

[3]

[4]

[5]

[6]

[7]

[8]

[9]

[10]

[11]

[12]