Llama 3, Mistral, Qwen3가 모두 Grouped-Query Attention(GQA)을 기본 아키텍처로 채택하면서, 모델 스펙에 적힌 num_kv_heads는 단순한 아키텍처 파라미터가 아니라 GPU 메모리 예산과 배치 크기 상한선을 결정하는 인프라 변수가 됐다. 모델을 고를 때 어텐션 헤드 수를 서빙 환경 제약과 함께 보지 않으면 운영에서 예상 못한 병목을 맞는다.
세 구조가 KV 헤드를 쌓는 방식
Multi-Head Attention(MHA)은 쿼리, 키, 밸류 헤드를 각각 H개 유지한다. Shazeer(2019)가 제안한 Multi-Query Attention(MQA)은 키와 밸류를 단 1개 헤드로 줄이고 쿼리만 H개 유지한다. 목적은 decode 단계에서 KV 캐시를 읽는 HBM 대역폭 소모를 줄이는 것이었다. Ainslie et al.(2023)의 Grouped-Query Attention(GQA)은 그 중간이다 — 키·밸류 헤드를 G개 그룹으로 묶어, 각 그룹 내 쿼리들이 하나의 K·V를 공유한다. G=1이면 MQA, G=H이면 MHA와 같다.
서빙에서 이 차이가 만드는 결과는 하나다. KV 캐시 크기가 num_kv_heads에 선형 비례한다.
KV 캐시 메모리: 수식으로 직접 계산
KV cache bytes = 2 × num_kv_heads × head_dim × seq_len × num_layers × bytes_per_element
2는 K와 V를 각각 저장하기 때문이다. 다른 파라미터가 같으면 num_kv_heads가 유일한 차이를 만든다.
Llama 3 8B를 대입해보자. 이 모델은 32 query heads, 8 KV heads(GQA-8), head_dim 128, 32 layers, bf16(2 bytes)다. 같은 모델이 MHA(num_kv_heads=32)나 MQA(num_kv_heads=1)를 썼다면 KV 캐시가 어떻게 달라지는지 비교한다.
요청 1건, GQA-8 기준 계산:
2 × 8 × 128 × seq_len × 32 × 2 = 131,072 × seq_len bytes
| seq_len | MHA (32 KV heads) | GQA-8 (8 KV heads) | MQA (1 KV head) |
|---|---|---|---|
| 2,048 | 1.0 GB | 0.25 GB | 0.031 GB |
| 8,192 | 4.0 GB | 1.0 GB | 0.125 GB |
| 32,768 | 16.0 GB | 4.0 GB | 0.5 GB |
| 131,072 | 64.0 GB | 16.0 GB | 2.0 GB |
seq_len 131,072에서 MHA 단일 요청의 KV 캐시 하나가 A100 80GB를 통째로 채운다. GQA-8은 같은 조건에서 16 GB다. 이 수치 자체보다 중요한 건 배치로 곱해질 때다.
KV 캐시 절감이 배치 크기 상한선을 올리는 과정
Llama 3 8B bf16 모델 가중치는 약 16 GB를 차지한다. A100 80GB에서 가중치와 활성화, 시스템 오버헤드를 제외하면 KV 캐시에 쓸 수 있는 공간은 대략 60 GB다. seq_len 8,192 기준으로 계산하면:
| 구조 | 요청당 KV 캐시 | 최대 동시 요청 (60 GB 기준) |
|---|---|---|
| MHA (32 KV heads) | 4.0 GB | ~15 |
| GQA-8 (8 KV heads) | 1.0 GB | ~60 |
| MQA (1 KV head) | 0.125 GB | ~480 |
GQA-8은 MHA 대비 4배 많은 요청을 동시에 받을 수 있다. 이 배치 크기 확장이 decode throughput으로 이어지는 경로는 직접적이다.
Decode는 token 1개를 생성할 때마다 이전 모든 token의 K·V 캐시 전체를 HBM에서 읽는 memory-bound 연산이다. 같은 HBM 대역폭 아래 배치가 클수록 요청당 연산량이 늘어나 arithmetic intensity가 높아지고, GPU core가 놀지 않는 시간이 줄어든다. KV 캐시 절감 → 배치 확장 → arithmetic intensity 상승 → tokens/s 개선은 하나로 연결된 연쇄다.
Ainslie et al.(2023)의 T5-XXL 기준 inference time 실측: MHA 1.51초, MQA 0.24초(약 6.3배 빠름), GQA-8 0.28초(약 5.4배 빠름). GQA-8과 MQA의 차이는 0.04초로 거의 없다.
TTFT·TPOT에 미치는 영향
Prefill과 Decode에서 GQA가 주는 이점의 성격이 다르다.
Prefill에서 GQA는 attention FLOPs를 줄인다. num_kv_heads가 32에서 8로 줄면 K·V projection과 attention score 계산 FLOPs가 그만큼 감소한다. 그런데 Prefill은 compute-bound 단계라, 전체 시간에서 FFN이 attention보다 훨씬 큰 비중을 차지한다. 실측에서 TTFT 개선이 KV 캐시 절감 비율만큼 선형으로 나오지 않는 이유가 여기 있다. GQA로 전환한다고 TTFT가 4배 줄지는 않는다.
Decode에서는 다르다. Token 1개를 생성할 때마다 이전 모든 token의 K·V를 HBM에서 읽어야 한다. KV 헤드가 줄면 이 로드 데이터 양이 직접 줄고, HBM 대역폭 병목이 완화된다. A100의 HBM 대역폭은 2 TB/s인데, 배치가 크고 시퀀스가 길수록 KV 캐시 로드가 병목이 된다. TPOT(Time Per Output Token)는 decode 속도를 직접 측정하는 지표라 GQA의 개선이 여기서 뚜렷하게 나타나고, 큰 배치일수록 그 효과가 더 크다.
vLLM은 num_kv_heads < num_query_heads이면 FlashAttention-2의 GQA 최적화 경로를 선택한다. FA2부터 grouped K·V 처리를 지원해, K·V 텐서를 query head 수에 맞게 repeat하지 않고 묶음 단위로 attention을 계산한다. 하드웨어에 따라 FA2(기본), FA3(Hopper H100/H200), FA4(Blackwell B200)가 선택된다.
그룹 수 선택의 트레이드오프
Ainslie et al.(2023)의 T5-XXL 기준 평균 품질 점수: MHA 47.2, GQA-8 47.1, MQA 46.6. GQA-8은 MHA 대비 0.1점 차이지만, MQA는 0.6점 낮다. 수치만 보면 차이가 작지만 long-context 태스크에서는 양상이 달라진다.
MQA가 K·V 헤드를 1개로 줄이면 각 레이어가 서로 다른 어텐션 패턴을 포착할 여지가 크게 줄어든다. 짧은 시퀀스에서는 티가 안 나지만, 수만 토큰 이상의 컨텍스트에서 K·V 표현력 부족이 검색·요약 품질에 누적된다. Llama 3가 128k 컨텍스트를 지원하면서 MQA가 아닌 GQA-8을 선택한 배경이다. MHA 대비 KV 캐시를 4분의 1로 줄이면서 품질 손실은 0.1점으로 잡았다.
GQA 그룹 수를 고를 때 실무에서 놓치기 쉬운 제약이 하나 더 있다. Tensor Parallel 설정과 kv_heads 수가 맞아야 한다. kv_heads는 TP GPU 수의 배수여야 한다. kv_heads=8 모델을 TP=8로 올리면 GPU당 head 1개로 균등 분할돼 문제없고, TP=4나 TP=2도 된다. 반면 kv_heads=4인 모델을 TP=8에 올리면 4를 8로 균등 분할할 수 없어 런타임 오류가 발생한다. 인프라팀에서 GPU 수를 늘리거나 TP 설정을 조정할 때 뒤늦게 이 제약에 걸리는 경우가 잦다. 모델을 선택하는 단계에서 kv_heads와 TP 설정을 함께 확인해야 한다.