GQA (Grouped-Query Attention) 분석

GQA — MHA와 MQA 사이의 KV 헤드 설계

Grouped-Query Attention(GQA)은 Query head는 여러 개 유지하면서 Key/Value head만 줄이고, 여러 Query head가 하나의 KV head를 공유하는 Attention 구조입니다. Multi-Head Attention(MHA)의 표현력과 Multi-Query Attention(MQA)의 작은 KV cache 사이에 중간 지점을 둡니다.

GQA의 직접적인 목적은 모델 전체 파라미터를 크게 줄이는 것이 아니라 autoregressive decoding에서 반복해서 읽는 KV cache의 용량과 메모리 대역폭을 낮추는 것입니다. Ainslie et al.은 기존 MHA checkpoint를 그룹화된 구조로 바꾼 뒤 원래 사전학습 계산량의 약 5%를 사용하는 uptraining recipe를 제시했고, GQA가 MQA에 가까운 속도와 MHA에 가까운 품질을 제공할 수 있음을 보였습니다.

그림 1. Query head 수는 유지하고 KV head 수만 H에서 G로 줄이는 구조

핵심 개념

Query head와 KV head의 분리

Query head 수를 H, Key/Value head 수를 G라고 하면 세 구조는 다음처럼 표현할 수 있습니다.

Hugging Face Transformers의 Llama 계열 설정에서는 num_attention_heads=H, num_key_value_heads=G로 이 구분을 나타냅니다. G=H이면 MHA, G=1이면 MQA, 그 사이면 GQA입니다. 일반적인 구현은 H가 G로 나누어떨어지도록 모델을 구성하여 각 KV head가 같은 수의 Query head를 담당하게 합니다.

KV cache 용량

레이어 수를 L, 배치를 B, 문맥 길이를 T, head dimension을 D, KV 원소 하나의 바이트 수를 s라고 하면 KV cache의 구조적 용량은 다음과 같이 근사할 수 있습니다.

KV_bytes ≈ 2 × L × B × T × G × D × s

앞의 2는 Key와 Value를 함께 저장한다는 뜻입니다. 같은 L, B, T, D, 자료형을 유지하면 GQA의 cache는 MHA 대비 G/H 비율입니다. 예를 들어 H=32, G=8인 GQA-8은 head 차원만으로 KV cache를 1/4로 줄입니다. 실제 사용량에는 block allocator, padding, scale metadata, tensor-parallel shard가 추가되므로 이 식은 설계 비교용으로 사용해야 합니다.

Grouped broadcast

입력 X에서 Query는 H개 head로, Key와 Value는 G개 head로 투영합니다.

Q = X · Wq  -> [B, H, T, D]
K = X · Wk  -> [B, G, T, D]
V = X · Wv  -> [B, G, T, D]

각 Query head h는 다음 그룹의 KV를 사용합니다.

g(h) = floor(h / (H / G))
score[h] = softmax(Q[h] · K[g(h)]ᵀ / sqrt(D))
O[h]     = score[h] · V[g(h)]

논리적으로는 K/V를 H/G번 반복한 것과 같지만, 효율적인 attention kernel은 동일한 K/V를 실제로 복제하지 않고 group index로 참조할 수 있습니다. 따라서 GQA는 하나의 attention 결과를 모든 head에 복사하는 방식이 아닙니다. Query마다 다른 score와 output을 계산하되, 읽는 K/V sequence만 그룹 내부에서 공유합니다.

그림 2. 각 Query group이 하나의 KV sequence를 공유하는 attention 경로

비교/분석

MHA, GQA, MQA 비교

GQA 그룹 수 선택

G를 작게 하면 cache와 KV read volume은 줄지만 한 KV 표현을 공유하는 Query head 수가 늘어 품질 저하 위험이 커질 수 있습니다. 반대로 G를 크게 하면 MHA에 가까운 표현력을 유지하지만 메모리 절감 폭과 decode 대역폭 이득이 줄어듭니다. 따라서 G는 단순히 최대 절감값으로 고르는 것이 아니라 모델의 품질 평가와 GPU/tensor parallel 배치 제약을 함께 보고 정해야 합니다.

예를 들어 H=32, D=128, L=32, B=1, T=4096, BF16(s=2)이면 다음과 같습니다.

이 값은 한 요청의 모든 레이어와 4096개 토큰을 합친 이론값입니다. allocator metadata, batch, padding, 양자화 scale은 포함하지 않았습니다.

Prefill과 Decode의 차이

Prefill에서는 여러 입력 토큰을 병렬로 처리하므로 GEMM 처리량과 연산 효율이 중요합니다. Decode에서는 새 Query가 보통 한 토큰씩 들어오고, 각 레이어가 과거 T개 토큰의 K/V를 다시 읽습니다. 이때 GQA는 과거 sequence의 KV head 수 자체를 G개로 줄여 HBM read pressure를 낮춥니다.

GQA가 attention의 모든 연산을 G/H 비율로 줄이는 것은 아닙니다. Query projection, Query head별 score 계산, output projection은 여전히 H개 Query head를 기준으로 수행됩니다. 실제 end-to-end 이득은 context length, batch size, attention kernel, memory bandwidth, prefill/decode 비율에 따라 달라집니다.

동작 원리

1. MHA checkpoint에서 GQA로 변환

기존 MHA의 K/V projection에는 H개의 head가 있습니다. 목표 GQA가 G개의 KV head를 사용하고 H/G가 정수라면, 원래 K/V head를 그룹별로 묶고 각 그룹의 가중치를 평균 풀링하여 새 K/V projection을 초기화할 수 있습니다. Query projection과 output projection은 유지합니다.

그림 3. 그룹별 K/V 평균 풀링과 uptraining을 통한 checkpoint 변환

for g in range(G):
    start = g * (H // G)
    end = (g + 1) * (H // G)
    Wk_g = mean(Wk_heads[start:end], axis=head)
    Wv_g = mean(Wv_heads[start:end], axis=head)

평균 풀링은 변환 직후의 초기화 방법이며, 품질 회복을 보장하는 것 자체는 아닙니다. GQA 논문은 변환한 checkpoint를 제한된 추가 학습으로 보정하는 uptraining을 사용했습니다. 실제 전환에서는 tokenizer, RoPE 설정, optimizer 상태, calibration 데이터와 평가 workload를 모델별로 맞춰야 합니다.

2. Projection과 RoPE

각 decoder layer는 hidden state에서 Q, K, V를 생성합니다. RoPE를 사용하는 모델이라면 Query와 Key에 위치 회전을 적용한 뒤 Key를 KV cache에 추가합니다. GQA는 Key/Value head 수를 줄이는 구조이므로 Query head별 위치 처리와 score 계산은 독립적으로 남습니다.

Q_t = RoPE(project_Q(x_t))       # [B, H, 1, D]
K_t = RoPE(project_K(x_t))       # [B, G, 1, D]
V_t = project_V(x_t)              # [B, G, 1, D]
append(K_cache, K_t)
append(V_cache, V_t)

3. Incremental decode

새 토큰 t를 처리할 때 현재 Query Q_t[h]는 자신이 속한 그룹 g(h)의 전체 cache와 내적합니다. 그룹이 같아도 Query가 다르므로 attention probability와 output은 서로 다릅니다.

for each layer:
    q_t = project_query(x_t)       # H heads
    k_t = project_key(x_t)         # G heads
    v_t = project_value(x_t)       # G heads
    append(K_cache, k_t)
    append(V_cache, v_t)

    for h in range(H):
        g = h // (H // G)
        a_h = softmax(q_t[h] @ K_cache[g].T / sqrt(D))
        o_h = a_h @ V_cache[g]

    y_t = concat(o_0, ..., o_{H-1}) @ Wo

이 경로에서 cache의 시간축 길이 T는 동일하지만 KV head 축이 H에서 G로 줄어듭니다. 커널이 K/V를 materialize하지 않고 group mapping을 처리하면 중간 broadcast buffer도 줄일 수 있습니다.

4. Tensor parallelism과 kernel 제약

Tensor parallelism에서는 Query head와 KV head가 GPU shard에 어떻게 배치되는지가 중요합니다. G가 GPU 수 또는 tensor-parallel degree와 잘 맞지 않으면 KV head를 균등하게 나누기 어렵고, attention 전에 KV를 복제하거나 통신하는 경로가 생길 수 있습니다. 따라서 이론적인 G/H cache 감소가 같은 비율의 latency 감소로 이어진다고 가정하면 안 됩니다.

구현 시 다음을 확인해야 합니다.

• num_attention_heads, num_key_value_heads, head_dim의 정합성과 H % G == 0 조건
• KV cache layout이 backend가 요구하는 [batch, kv_heads, sequence, head_dim] 계열인지 여부
• repeat_kv가 실제 메모리 복제인지 fused kernel 내부의 논리적 broadcast인지 여부
• tensor parallel shard별 KV head 배치와 all-gather/broadcast 통신량
• GQA와 FlashAttention, PagedAttention, KV cache quantization의 shape·scale 호환성

장단점

장점

• MHA 대비 KV cache 용량과 decode 단계의 KV memory read를 G/H 수준으로 줄일 수 있습니다.
• Query head 수를 유지하므로 MQA보다 표현 공간과 품질을 더 보존할 여지가 있습니다.
• 긴 context, 큰 batch, beam search처럼 KV cache가 커지는 serving workload에 유리합니다.
• G를 조절해 품질·메모리·대역폭 사이의 절충점을 모델별로 선택할 수 있습니다.
• PagedAttention은 cache 배치 방식을, KV quantization은 원소 정밀도를 바꾸므로 GQA와 직교적으로 결합할 수 있습니다.

단점

• K/V 표현을 공유하므로 MHA의 head별 독립성보다 표현력이 줄고 품질 저하가 발생할 수 있습니다.
• 기존 MHA checkpoint를 단순히 head를 삭제하는 방식으로 변환하면 분포가 바뀌므로 평균 풀링과 추가 학습이 필요할 수 있습니다.
• 작은 G는 tensor parallel shard에 KV head를 배치하기 어렵게 만들어 broadcast 또는 통신 비용을 유발할 수 있습니다.
• Prefill이 compute-bound이면 decode에서의 KV read 절감이 전체 요청 latency에 제한적으로만 반영됩니다.
• 모델 설정을 바꾸는 아키텍처 기법이므로 기존 weight와 inference backend가 GQA shape을 지원해야 합니다.

관련 기술 및 참고 문헌

핵심 정리

GQA는 H개의 Query head를 유지하면서 G개의 KV head만 두고, 각 KV head를 H/G개의 Query head가 공유하는 Attention 구조입니다. KV cache 용량은 KV head 수에 비례하므로 MHA 대비 G/H 수준으로 줄지만, Query 계산과 모든 head의 score 계산이 같은 비율로 줄어드는 것은 아닙니다. 기존 MHA checkpoint는 그룹별 K/V 평균 풀링 후 제한된 uptraining으로 GQA에 맞출 수 있으며, G 선택은 품질뿐 아니라 tensor parallel 배치와 kernel 지원까지 고려해야 합니다. PagedAttention과 KV cache quantization은 서로 다른 축을 최적화하므로 GQA와 함께 적용할 수 있습니다.