MQA (Multi-Query Attention) 분석

MQA (Multi-Query Attention) — 공유 KV 헤드와 디코딩 메모리 대역폭 절감

Multi-Query Attention(MQA)은 여러 Query head가 하나의 Key head와 Value head를 공유하는 Attention 변형입니다. Query head는 그대로 유지하므로 여러 표현 공간에서 점수를 계산할 수 있지만, autoregressive decoding에서 반복해서 읽는 KV cache는 훨씬 작아집니다. 따라서 MQA의 주된 목적은 학습 FLOPs를 줄이는 것이 아니라, 한 토큰씩 생성할 때 발생하는 KV cache 메모리 대역폭과 용량을 줄이는 것입니다.

MQA는 2019년 Shazeer의 Fast Transformer Decoding: One Write-Head is All You Need에서 제안되었습니다. 원 논문은 WMT14 English-German 실험에서 multi-head baseline의 decoder step 비용 46 µs가 MQA에서 3.8 µs로 줄어드는 결과를 보고했지만, 이는 TPUv2와 고정된 실험 조건의 측정값입니다. 이후 GQA는 Query head와 KV head 사이에 중간 개수의 그룹을 두어 MQA의 속도와 MHA의 품질 사이를 절충했습니다.

핵심 개념

MHA, MQA, GQA의 head 구성

일반적인 Multi-Head Attention(MHA)은 각 Query head에 별도의 K/V head를 둡니다. MQA는 모든 Query head가 하나의 K/V head를 함께 사용합니다. Grouped-Query Attention(GQA)은 Query head를 여러 그룹으로 묶고, 그룹마다 하나의 K/V head를 배치합니다.

그림 1. Query head 수는 유지하면서 KV head 수만 줄이는 MQA/GQA의 구조 비교

Hugging Face Transformers의 Llama 계열 설정도 num_key_value_heads로 이 구분을 표현합니다.

• num_key_value_heads = num_attention_heads: MHA
• num_key_value_heads = 1: MQA
• 1 < num_key_value_heads < num_attention_heads: GQA

따라서 MQA는 Query까지 하나로 합치는 방식이 아닙니다. Query는 head별로 독립적으로 생성되며, K/V만 공유됩니다.

KV cache의 크기

레이어 수를 L, Query head 수를 H, KV head 수를 G, head dimension을 D, 시퀀스 길이를 T, 배치를 B, KV 자료형의 바이트 수를 s라고 하면, 단순한 KV cache 용량은 다음과 같이 근사할 수 있습니다.

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

앞의 2는 Key와 Value를 함께 저장하기 때문입니다. MHA에서는 G=H, MQA에서는 G=1, GQA에서는 1<G<H입니다. 같은 H, D, L, B, T, 자료형이라면 MQA의 KV 용량은 MHA의 1/H, GQA는 G/H 수준입니다. 실제 구현에서는 tensor parallelism의 shard, padding, block allocator, scale 메타데이터가 추가되므로 이 식은 구조적인 비교용으로 보는 것이 안전합니다.

왜 decoding에서 특히 효과적인가

Prefill은 여러 위치의 Query를 병렬로 처리할 수 있어 GEMM의 계산량과 GPU 활용률이 중요합니다. 반면 decode는 보통 새 Query 한 개만 추가하고 과거 전체 K/V를 다시 읽습니다. 이때 연산량보다 HBM에서 KV를 가져오는 비용이 병목이 되기 쉬우므로, KV head 수를 줄이는 MQA의 효과가 커집니다.

MQA는 다음 네 가지를 동시에 줄입니다.

• KV cache의 저장 용량
• 매 decoding step에서 읽어야 하는 K/V 데이터량
• K/V projection 결과의 저장·이동량
• 대규모 배치에서 요청 하나가 차지하는 메모리 발자국

반면 Query projection, attention score 계산의 Query head 수, output projection은 기본적으로 유지됩니다. 따라서 MQA를 "Attention 전체가 1개 head가 되는 기법"으로 이해하면 안 됩니다.

비교/분석

간단한 예

H=32, D=128, FP16(s=2)인 한 레이어에서 토큰 하나의 K/V 저장량은 다음과 같습니다.

이는 한 레이어·한 토큰만의 값입니다. 32개 레이어, 긴 문맥, 여러 요청이 겹치면 차이가 누적됩니다. 다만 KV head를 줄였다고 해서 Query head에 필요한 연산과 모델 가중치 전체가 같은 비율로 줄어드는 것은 아닙니다.

정확도와 속도의 절충

KV를 공유하면 각 head가 서로 다른 K/V projection을 사용하던 자유도가 사라집니다. 원 논문은 MQA가 baseline보다 약간 낮은 품질을 보이지만 대체로 비슷한 수준이라고 보고했습니다. 이후 GQA 논문은 기존 MHA checkpoint를 평균 pooling 등으로 KV head를 줄인 뒤 원래 pre-training compute의 5%를 사용하는 uptraining으로 품질을 회복하는 방법을 제시했고, 중간 개수의 KV head가 MQA에 가까운 속도와 MHA에 가까운 품질을 제공한다고 보고했습니다.

동작 원리

1. Projection 단계

입력 hidden state X에서 Query는 H개 head로 투영하고, Key와 Value는 G개 head로 투영합니다. MQA에서는 G=1이므로 K/V projection의 출력 shape가 Query projection보다 작습니다.

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

MHA의 경우 G=H이고, MQA의 경우 G=1입니다. GQA에서는 각 Query head가 H/G개의 Query group 중 하나에 속합니다. H가 G로 나누어 떨어지는 설정이 일반적이지만, 실제 프레임워크는 모델 설정의 제약을 별도로 검사합니다.

2. Attention에서 K/V broadcast

각 Query head는 자신이 속한 KV group의 K와 V를 사용합니다. 논리적으로는 K/V를 Query head 수에 맞춰 broadcast 또는 repeat한 것과 같지만, 효율적인 커널은 실제로 동일한 K/V를 매번 복제하지 않고 group index로 참조할 수 있습니다.

그림 2. 한 토큰 decode에서 Query는 H개를 유지하고 K/V cache는 G개만 읽는 흐름

group(h) = floor(h / (H / G))
score[h] = softmax(Q[h] · K[group(h)]^T / sqrt(D))
O[h]     = score[h] · V[group(h)]
Y        = concat(O[0], ..., O[H-1]) · Wo

MQA에서는 모든 h에 대해 group(h)=0입니다. 즉, score vector는 head별로 다르지만 모든 head가 같은 K/V sequence를 조회합니다. 이 차이 때문에 MQA는 하나의 Attention 결과를 모든 head에 복사하는 것이 아닙니다.

3. Incremental decoding과 cache append

새 토큰의 hidden state x_t가 들어오면 각 레이어는 새 q_t를 H개 만들고, 새 k_t, v_t는 G개만 만듭니다. 새 K/V는 cache의 마지막 위치에 추가되고, 현재 Query는 이전 모든 위치의 K/V와 attention을 수행합니다.

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 query_heads:
        g = h // (H / G)
        a_h = softmax(q_t[h] @ K_cache[g].T / sqrt(D))
        o_h = a_h @ V_cache[g]

이 경로에서 매 step 읽는 cache의 길이 T는 같지만 head 수 G가 줄어듭니다. 그래서 긴 문맥이나 큰 batch에서 HBM read pressure가 줄고, 같은 GPU 메모리로 더 많은 sequence를 수용할 수 있습니다.

4. Tensor parallelism 주의점

여러 GPU에 모델을 tensor parallel로 나눌 때는 Query head와 KV head의 배치가 통신량에 영향을 줍니다. KV head가 너무 적으면 각 shard에 KV head를 균등하게 배치하기 어렵거나, attention 전에 KV를 broadcast하는 통신이 필요할 수 있습니다. 따라서 이론적인 H/G 메모리 절감이 그대로 end-to-end latency 향상으로 이어지는지는 GPU 수, head 배치, fused attention kernel, interconnect에 따라 달라집니다.

장단점

장점

• KV head 수를 줄여 KV cache 용량과 매 decode step의 메모리 읽기량을 직접 줄입니다.
• Query head 수를 유지하므로 단순히 전체 Attention을 저차원으로 축소하는 것보다 표현력을 보존합니다.
• 긴 context, 큰 batch, beam search와 같이 KV cache가 커지는 serving workload에 특히 유리합니다.
• GQA라는 중간 지점을 통해 품질과 메모리 대역폭 사이를 모델별로 조정할 수 있습니다.
• PagedAttention, continuous batching, KV cache quantization과 직교적으로 결합할 수 있습니다. MQA는 head 차원, PagedAttention은 배치·토큰의 배치 방식, quantization은 원소 정밀도를 줄이는 기법입니다.

단점

• 모든 Query head가 같은 K/V 표현을 공유하므로 MHA 대비 품질 저하가 생길 수 있습니다.
• 기존 MHA checkpoint를 단순히 잘라서 MQA로 바꾸면 분포가 달라집니다. mean pooling 기반 uptraining이나 별도 재학습이 필요할 수 있습니다.
• KV head가 1개라서 GPU·tensor parallel shard에 균등하게 배치하기 어렵고, 구현에 따라 broadcast 통신이 생길 수 있습니다.
• Prefill에서는 decode만큼 KV read bandwidth가 지배적이지 않을 수 있어, MQA의 효과가 workload 전체에서 동일하지 않습니다.
• Query 계산과 output projection은 크게 줄지 않으므로, MQA를 적용해도 compute-bound workload의 속도가 같은 비율로 개선되지는 않습니다.

관련 기술 및 참고 문헌

구현 시 확인할 항목

• 모델 config의 num_attention_heads, num_key_value_heads, head_dim이 서로 일관적인지 확인합니다.
• KV cache shape가 [batch, kv_heads, sequence, head_dim]인지, backend가 요구하는 layout인지 확인합니다.
• GQA/MQA의 KV repeat가 실제 복제인지, attention kernel 내부의 논리적 broadcast인지 확인합니다.
• MQA와 quantized cache를 함께 사용할 때 scale의 축과 KV group mapping이 일치하는지 확인합니다.
• 품질은 perplexity만 보지 말고 long-context retrieval, instruction following, reasoning 등 실제 workload에서 측정합니다.

핵심 정리

MQA는 H개의 Query head를 유지하면서 K/V head를 1개로 공유해 KV cache의 용량과 decoding 메모리 대역폭을 줄이는 Attention 구조입니다. KV cache 용량은 KV head 수 G에 비례하므로 MHA 대비 GQA는 G/H, MQA는 1/H 수준으로 줄어들 수 있습니다. 대신 K/V 표현을 공유하는 대가로 품질 저하와 tensor parallel 배치 제약이 생길 수 있어, 기존 checkpoint에는 uptraining이나 GQA가 실용적인 절충이 됩니다. PagedAttention과 quantization이 각각 배치 방식과 원소 정밀도를 줄이는 것과 달리, MQA는 Attention의 head 차원을 줄이는 기법이며 세 가지를 함께 적용할 수 있습니다.

주의 — 원 논문의 46 µs→3.8 µs 수치는 WMT14 English-German, TPUv2, 특정 batch와 sequence length의 incremental decoding 측정값입니다. 실제 이득은 모델의 H/G, context length, batch size, attention kernel, tensor parallelism, interconnect와 prefill/decode 비율에 따라 달라집니다.