MQA와 GQA: KV 캐시 경량화 전략

지난 글에서 어텐션의 $O(N^2)$ 비용을 줄이는 두 갈래를 봤다. 참조를 잘라 내는 희소 어텐션·슬라이딩 윈도·SSM과, 결과를 그대로 둔 채 GPU 메모리 이동만 줄이는 FlashAttention이다. 둘 다 어텐션을 계산하는 비용에 손대고, 계산이 끝난 뒤 다음 토큰을 위해 남겨 두는 값에는 손대지 않는다. 이번에는 그 남겨 두는 값, 곧 KV 캐시를 줄이는 두 가지 기법인 MQA(Multi-Query Attention)와 GQA(Grouped-Query Attention)를 다룬다. LLaMA 3, Mistral, Qwen2처럼 2023년 이후 나온 오픈 LLM 대부분이 GQA를 쓴다.

두 기법은 어텐션의 계산 방식을 거의 바꾸지 않는다. 바꾸는 것은 K와 V를 몇 벌 만들어 두느냐 하나다. 그런데 그 한 가지가 긴 문맥과 큰 배치에서 GPU 한 장이 몇 명의 요청을 동시에 받을 수 있는지를 정한다. 이 글은 그 용량을 식으로 세워 직접 계산해 보는 데서 시작해, 헤드를 공유하면 무엇을 잃는지, 이미 학습한 모델을 어떻게 바꾸는지, 그리고 GQA 다음에 나온 압축 방식과 서빙 쪽 기법이 이것과 어떻게 맞물리는지까지 따라간다.

KV 캐시

캐시가 하는 일

LLM은 문장을 한 토큰씩 만든다. 앞 토큰들을 보고 다음 토큰 하나를 고르고, 그 토큰을 붙여 다시 다음을 고르는 자기회귀 디코딩이다. 이때 어텐션은 새 토큰의 쿼리(Q)를 앞선 모든 토큰의 키(K)와 비교하고, 그 비율대로 값(V)을 섞는다. 앞 토큰들의 K와 V는 한 번 계산하면 다시 바뀌지 않으므로, 매 걸음마다 새로 계산하지 않고 GPU 메모리에 쌓아 두고 꺼내 쓴다. 이렇게 쌓아 두는 저장소가 KV 캐시다.

캐시가 없으면 100번째 토큰을 만들 때 앞 99개 토큰의 K·V를 레이어마다 다시 투영해야 하고, 1,000번째 토큰에서는 999개를 다시 투영한다. 캐시가 있으면 새 토큰 하나의 K·V만 계산해 끝에 붙이면 된다. 계산을 메모리와 맞바꾼 셈이고, 그 메모리가 문맥 길이와 동시 요청 수에 정비례해서 불어난다는 것이 이 글 전체의 출발점이다.

용량 공식

KV 캐시의 크기는 곱셈 하나로 나온다.

$$\text{KV 캐시} = 2 \times n_{kv} \times d_h \times L \times S \times B \times \text{바이트}$$

맨 앞의 2는 K와 V 두 벌이라는 뜻이고, $n_{kv}$ 는 KV 헤드 수, $d_h$ 는 헤드 하나의 차원, $L$ 은 레이어 수, $S$ 는 문맥 길이, $B$ 는 동시에 처리하는 시퀀스 수(배치), 마지막은 원소 하나의 바이트 수(bf16이면 2)다. 여기서 모델 설계가 정하는 것은 앞의 셋과 레이어 수이고, 문맥 길이와 배치는 서비스가 정한다. 이 글의 두 기법이 건드리는 것은 $n_{kv}$ 하나뿐이다.

LLaMA 3 70B에 숫자를 넣어 보자. Q 헤드 64개, KV 헤드 8개, 헤드 차원 128, 레이어 80개다.

토큰 하나당 (GQA, KV 헤드 8)
  2 × 8 × 128 × 80 × 2바이트 = 327,680바이트 ≈ 0.33 MB

문맥 8,192 · 배치 1  →  약 2.7 GB
문맥 8,192 · 배치 32 →  약 86 GB
문맥 128K  · 배치 1  →  약 43 GB

같은 모델이 KV 헤드를 64개 모두 가졌다면(MHA) 토큰당 2.6 MB이고 8K 문맥 한 줄에 21.5 GB다. 배치 32면 약 690 GB로, bf16 가중치 140 GB의 다섯 배가 KV 캐시 하나에 들어간다. 모델을 올리는 데 드는 메모리보다 사용자를 받는 데 드는 메모리가 훨씬 크다는 것, 그리고 그 차이가 KV 헤드 수에 정비례한다는 것이 공식이 알려 주는 핵심이다.

이 공식을 뒤집으면 서비스 설계에 바로 쓰는 질문이 된다 — 이 GPU에 동시 요청을 몇 개 받을 수 있는가. 80 GB짜리 GPU 네 장에 70B 모델을 bf16으로 올리면 가중치가 140 GB를 먹고 180 GB가 남는다. 활성값과 런타임이 쓰는 몫으로 20 GB쯤 떼면 KV 캐시에 쓸 수 있는 것은 약 160 GB다. GQA라면 8K 문맥 요청을 60개 가까이 동시에 태울 수 있고, MHA였다면 일곱 개가 한계다. 같은 하드웨어에서 받을 수 있는 사용자 수가 여덟 배 차이 나는 것이고, 요청당 GPU 비용도 거의 그 비율로 갈린다.

대역폭 병목

메모리 크기만 문제가 아니다. 디코딩 한 걸음은 토큰 하나를 만들기 위해 가중치 전부와 그 시퀀스의 KV 캐시 전부를 GPU 메모리에서 읽어 와야 한다. 곱셈 자체는 적은데 읽을 양이 많아서, 이 단계의 속도는 연산 능력이 아니라 메모리 대역폭이 정한다.

배치를 키우면 가중치를 한 번 읽어 여러 시퀀스에 나눠 쓰므로 가중치 읽기 비용이 분산된다. 하지만 KV 캐시는 시퀀스마다 제 것이라 나눠 쓸 수 없다. 위의 MHA 예에서 배치 32의 한 걸음은 가중치 140 GB에 KV 690 GB를 읽는다 — 읽는 양의 80% 이상이 KV다. GQA로 KV를 86 GB로 줄이면 한 걸음에 읽는 총량이 830 GB에서 226 GB로 줄고, 그만큼 걸음이 빨라진다. KV 헤드를 줄이는 기법이 메모리 절감과 함께 디코딩 속도까지 올리는 것은 이 때문이다.

반대로 프롬프트 전체를 한 번에 읽어 들이는 첫 단계, 곧 프리필은 토큰 수천 개를 한꺼번에 행렬 곱으로 처리하므로 대역폭보다 연산이 병목이다. 그래서 KV 헤드를 줄여도 첫 토큰이 나오기까지의 시간은 크게 줄지 않고, 줄어드는 것은 그 뒤 토큰이 하나씩 이어 나오는 간격이다. 긴 답을 생성하는 서비스일수록 이 차이가 체감된다.

MQA

공유 K·V

2019년 Noam Shazeer가 제안한 MQA는 발상이 단순하다. Q는 헤드마다 따로 두고, K와 V는 한 벌만 만들어 모든 Q 헤드가 같이 쓴다. 헤드가 여럿인 이유는 서로 다른 관계를 따로 보기 위해서인데, MQA는 「무엇을 물을지」는 헤드마다 다르게 두고 「무엇을 꺼낼 수 있는지」는 하나로 묶은 셈이다.

class MultiQueryAttention(nn.Module):
    def __init__(self, d_model, num_heads, head_dim):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        # Q: 헤드별로 분리
        self.W_q = nn.Linear(d_model, num_heads * head_dim)
        # K, V: 단 1개 헤드 (모든 Q 헤드가 공유)
        self.W_k = nn.Linear(d_model, head_dim)
        self.W_v = nn.Linear(d_model, head_dim)
        self.W_o = nn.Linear(num_heads * head_dim, d_model)

    def forward(self, x, mask=None):
        B, T, _ = x.shape
        q = self.W_q(x).view(B, T, self.num_heads, self.head_dim)
        k = self.W_k(x).view(B, T, 1, self.head_dim)
        v = self.W_v(x).view(B, T, 1, self.head_dim)
        # k, v를 num_heads 차원으로 브로드캐스트
        k = k.expand(-1, -1, self.num_heads, -1)
        v = v.expand(-1, -1, self.num_heads, -1)
        # ... 이하 표준 어텐션

용량 공식의 $n_{kv}$ 가 1이 되므로 KV 캐시는 헤드 수만큼, 64헤드 모델이면 64분의 1로 준다. 위 70B 설정이라면 8K 문맥 한 줄이 0.34 GB다. PaLM, Falcon 7B, StarCoder가 이 방식을 썼다. 코드에서 expand는 실제로 복사하지 않고 같은 메모리를 여러 번 가리키게만 하므로, 추론 커널도 K·V 한 벌을 읽어 모든 헤드에 나눠 쓴다.

품질 저하

공짜는 아니다. 헤드마다 K·V를 따로 두면 한 헤드는 문법 관계에 맞춘 키 공간을, 다른 헤드는 지시어가 가리키는 대상에 맞춘 키 공간을 가질 수 있다. MQA에서는 모든 헤드가 같은 키 공간을 봐야 하므로, 쿼리를 아무리 다르게 만들어도 꺼낼 수 있는 정보의 모양이 한 가지로 제한된다. 표현력의 상한이 내려가는 것이다.

GQA를 제안한 Ainslie 등의 2023년 논문은 T5 모델로 요약·번역·질의응답을 재서 이 차이를 보였다. MQA로 바꾼 큰 모델은 같은 크기의 MHA보다 점수가 조금 낮았고, 대신 한 단계 작은 MHA보다는 높으면서 훨씬 빨랐다. 논문은 여기에 더해 MQA가 학습을 불안정하게 만들 수 있다는 점과, 추론을 빠르게 하려고 별도 모델을 처음부터 다시 학습하는 것 자체가 부담이라는 점을 동기로 들었다. 품질 차이는 평균 점수로는 작아 보이지만, 긴 입력에서 여러 곳에 흩어진 정보를 동시에 끌어와야 하는 작업처럼 헤드마다 다른 것을 봐야 하는 자리에서 먼저 드러난다고 보는 것이 일반적인 해석이다. 그래서 현장의 선택은 「조금 잃고 많이 얻을 것인가」가 되었고, 그 사이를 메운 것이 다음 절의 GQA다.

GQA

그룹 공유

GQA는 Q 헤드를 G개 그룹으로 나누고, 같은 그룹의 Q 헤드끼리만 KV 한 벌을 나눠 쓴다. G가 1이면 모든 헤드가 한 벌을 쓰니 MQA이고, G가 헤드 수와 같으면 헤드마다 한 벌이니 MHA다. 즉 GQA는 새로운 구조라기보다 MHA와 MQA를 양 끝으로 둔 한 줄의 눈금이고, 그 눈금 위 어디에 설지를 G로 고르는 방식이다.

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, num_heads, num_kv_heads, head_dim):
        super().__init__()
        assert num_heads % num_kv_heads == 0
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.groups = num_heads // num_kv_heads  # 그룹당 Q 헤드 수

        self.W_q = nn.Linear(d_model, num_heads * head_dim, bias=False)
        self.W_k = nn.Linear(d_model, num_kv_heads * head_dim, bias=False)
        self.W_v = nn.Linear(d_model, num_kv_heads * head_dim, bias=False)
        self.W_o = nn.Linear(num_heads * head_dim, d_model, bias=False)

    def forward(self, x, freqs_cos, freqs_sin):
        B, T, _ = x.shape
        q = self.W_q(x).view(B, T, self.num_heads, -1)
        k = self.W_k(x).view(B, T, self.num_kv_heads, -1)
        v = self.W_v(x).view(B, T, self.num_kv_heads, -1)
        # KV를 그룹 크기만큼 반복해 Q와 shape 맞춤
        k = k.repeat_interleave(self.groups, dim=2)
        v = v.repeat_interleave(self.groups, dim=2)
        # 이후 표준 어텐션 (FlashAttention 적용 가능)

이 코드의 repeat_interleave는 이해를 돕기 위한 것이고, 실제 추론 커널은 K·V를 복제하지 않는다. 캐시에는 8벌만 두고, 각 Q 헤드가 자기 그룹 번호로 그 벌을 찾아 읽는다. 그래서 절감은 계산이 아니라 저장과 읽기에서 온다.

그룹 수 고르기

LLaMA 3 70B(Q 헤드 64, 레이어 80, 헤드 차원 128)에 세 방식을 대 보면 절감 폭이 한눈에 보인다. 8K 문맥, 배치 1, bf16 기준이다.

방식	KV 헤드 수	KV 캐시	MHA 대비
MHA	64	약 21.5 GB	1
GQA	8	약 2.7 GB	1/8
MQA	1	약 0.34 GB	1/64

GQA 논문에서 그룹 8개짜리 모델은 품질이 MHA에 가깝고 속도는 MQA에 가까웠다. 그룹 수를 더 줄이면 KV는 줄지만 품질이 MQA 쪽으로 내려가고, 늘리면 그 반대다. 이 곡선이 초반에 가파르고 뒤로 갈수록 평평해서, 그룹 몇 개만 둬도 품질 대부분이 돌아온다. 공개 모델들이 KV 헤드 8개 근처에 몰려 있는 것은 이 모양 때문이다.

모델	Q 헤드	KV 헤드	그룹당 Q 헤드
LLaMA 3 8B	32	8	4
LLaMA 3 70B	64	8	8
Mistral 7B	32	8	4
Qwen2 72B	64	8	8
Gemma 2 9B	16	8	2

표를 보면 모델 크기가 달라도 KV 헤드는 8로 같고, 커질수록 그룹당 Q 헤드가 는다. 여기에는 서빙 쪽 이유도 있다. 큰 모델은 GPU 여러 장에 헤드를 나눠 싣는 텐서 병렬로 도는데, KV 헤드가 GPU 수로 나누어떨어져야 장마다 KV를 고르게 나눌 수 있다. GPU 8장짜리 서버가 흔하니 KV 헤드 8은 한 장에 한 벌씩 들어가는 수다. 8장보다 KV 헤드가 적으면 같은 KV를 여러 장에 복제해야 해서 절감 일부가 사라진다.

업사이클링

평균 풀링 초기화

이미 MHA로 학습한 모델을 GQA로 바꾸고 싶을 때 처음부터 다시 학습하는 것은 너무 비싸다. Ainslie 등은 기존 체크포인트를 고쳐 쓰는 업사이클링을 제안했다. 같은 그룹에 들어갈 헤드들의 K 투영 행렬과 V 투영 행렬을 각각 평균 내어 그룹의 KV 헤드 하나로 삼고, 그 상태로 짧게 추가 학습해 모델이 새 구조에 적응하게 한다.

def mha_to_gqa(w_k, num_heads, num_kv_heads):
    # w_k: (num_heads * head_dim, d_model)
    head_dim = w_k.shape[0] // num_heads
    w_k = w_k.view(num_heads, head_dim, -1)
    group_size = num_heads // num_kv_heads
    # 그룹별 평균으로 KV 헤드 초기화
    w_k_gqa = w_k.view(num_kv_heads, group_size, head_dim, -1).mean(dim=1)
    return w_k_gqa.view(num_kv_heads * head_dim, -1)

왜 평균인가. 논문은 그룹의 첫 헤드만 남기는 방법, 무작위로 새로 만드는 방법과 평균을 견줬고 평균이 가장 나았다. 평균은 그룹 안 헤드들이 각자 배운 것을 조금씩 담고 시작하지만, 하나만 남기면 나머지 헤드가 배운 것이 통째로 버려지고 무작위 초기화는 배운 것을 전부 버린다. 추가 학습이 메워야 할 거리가 평균에서 출발할 때 가장 짧은 것이다. 위 코드는 K만 보였지만 V 투영에도 같은 함수를 그대로 쓰고, Q 투영과 출력 투영은 손대지 않는다 — 헤드 수가 줄어드는 것은 K·V뿐이기 때문이다.

추가 학습량

추가 학습을 얼마나 하느냐가 실제 비용이다. 논문은 원래 사전학습 걸음 수의 5%만큼 더 학습한 모델을 주 결과로 냈다. 0%, 곧 평균만 내고 바로 쓰면 품질이 눈에 띄게 떨어졌고, 5%까지 올리면 대부분 회복됐으며, 10%로 늘려도 추가 이득은 작았다. 사전학습 전체를 다시 하는 것과 비교하면 스무 분의 1 비용으로 구조를 바꾸는 셈이다.

실무에서 이 숫자를 옮길 때는 조심할 점이 둘 있다. 첫째, 5%는 사전학습 데이터와 같은 분포로 학습했을 때의 값이다. 가진 데이터가 특정 도메인뿐이면 그 도메인에 치우친 모델이 나온다. 둘째, 평가를 사전학습 지표만으로 하면 안 된다. 긴 문맥 검색이나 여러 턴에 걸친 지시 따르기처럼 헤드 공유가 먼저 영향을 주는 작업을 따로 재 보고, 원본 MHA 모델과 나란히 비교한 뒤에 바꿔 넣는다. 처음부터 GQA로 학습하는 모델이 대부분인 지금은 이 절차가 주로 사내에서 오래 써 온 MHA 모델을 살려 쓸 때 쓰인다.

MLA

잠재 벡터 압축

GQA가 헤드 수를 줄여 캐시를 줄였다면, DeepSeek-V2(2024)가 도입한 MLA(Multi-head Latent Attention)는 다른 축을 줄인다. K와 V를 캐시에 두는 대신, 입력을 512차원짜리 작은 잠재 벡터로 압축해 그것만 캐시에 둔다. 어텐션을 계산할 때 그 잠재 벡터에서 헤드별 K와 V를 다시 펼쳐 쓴다.

이렇게 하면 헤드 수를 줄이지 않아도 된다. DeepSeek-V2는 헤드 128개를 그대로 두고도 토큰·레이어당 캐시가 576개 값이다. 같은 헤드 수의 MHA가 32,768개, KV 헤드 8개인 GQA가 2,048개이니 GQA보다도 작다. 논문은 이 크기가 그룹 2.25개짜리 GQA와 맞먹는다고 적고, 자사 이전 모델(DeepSeek 67B)보다 KV 캐시를 93.3% 줄였다고 보고했다.

펼치는 계산이 매번 붙으면 느려지지 않느냐는 질문이 자연스럽다. 여기에 요령이 있다. K를 펼치는 행렬은 Q 투영 행렬에 미리 곱해 합칠 수 있고, V를 펼치는 행렬은 출력 투영 행렬에 합칠 수 있다. 행렬 곱은 순서를 묶어 바꿀 수 있기 때문이다. 그러면 추론 때는 잠재 벡터를 펼치지 않고 바로 어텐션을 계산할 수 있다.

RoPE 분리

걸림돌이 하나 있었다. 요즘 LLM은 위치 정보를 K와 Q에 회전을 걸어 넣는 RoPE를 쓰는데, 이 회전은 토큰 위치마다 달라서 위에서 말한 행렬 합치기를 깨뜨린다. 위치마다 다른 회전 행렬이 사이에 끼면 두 행렬을 미리 곱해 둘 수가 없다.

DeepSeek-V2는 이 부분을 떼어 냈다. 위치에 반응해야 하는 몫은 64차원짜리 별도 키로 두고 모든 헤드가 함께 쓰게 했으며, 이 키에만 RoPE를 건다. 캐시에 남는 576은 이 64와 잠재 벡터 512를 더한 값이다. 헤드 공유라는 GQA의 발상이 여기서는 위치 키 한 조각에만 쓰였다는 점이 흥미롭다 — 위치 정보는 헤드마다 다르게 볼 필요가 적다는 판단이다. 이후 DeepSeek-V3도 이 구조를 이어 썼다. MLA는 구조를 바꾸는 일이라 이미 학습한 GQA 모델에 끼워 넣을 수 없고, 처음 설계할 때 고르는 선택지다.

지금까지 본 넷을 한 줄로 세우면 고르는 기준이 정리된다. MHA는 품질의 기준점이지만 긴 문맥과 큰 배치에서 캐시가 감당이 안 된다. MQA는 캐시를 가장 많이 줄이는 대신 품질과 학습 안정성에서 값을 치른다. GQA는 그 사이에서 KV 헤드 8개 근처를 고르면 품질을 거의 지키면서 캐시를 크게 줄이고, 기존 모델을 업사이클링으로 바꿀 수도 있어 가장 무난하다. MLA는 헤드 수를 지키면서 캐시를 GQA보다 더 줄이지만, 처음부터 그 구조로 학습해야 하고 서빙 엔진이 전용 커널을 갖춰야 제값을 한다. 새 모델을 설계하는 쪽이 아니라 공개 모델을 골라 서빙하는 쪽이라면, 이 비교는 「이 모델이 내 GPU에서 요청을 몇 개 받을 수 있는가」를 가늠하는 눈금으로 쓰면 된다.

함께 쓰는 절감

KV 캐시 양자화

헤드 수를 줄이는 것과 별개로, 캐시에 담는 원소 하나의 크기를 줄이는 방법이 있다. 용량 공식의 마지막 항인 바이트 수를 2에서 1로 내리면 FP8, 0.5로 내리면 4비트다. vLLM 같은 서빙 엔진은 FP8 KV 캐시를 설정 하나로 켜게 해 두었다.

두 절감은 곱해진다. GQA로 8분의 1, FP8로 다시 절반이면 MHA bf16 대비 16분의 1이다. 위 70B 예에서 8K 문맥 배치 32가 86 GB에서 43 GB로 줄어, 같은 GPU에 두 배의 요청을 올릴 수 있다.

대신 주의할 점이 있다. KV 헤드가 적을수록 헤드 하나가 더 많은 Q 헤드를 떠맡으므로, 그 헤드의 오차가 더 많은 곳으로 번진다. 또 K에는 특정 채널만 값이 크게 튀는 이상치가 있어서 토큰 단위로 눈금을 잡으면 그 채널 때문에 나머지 값이 뭉개진다. 4비트까지 내리는 연구(KIVI 등)가 K는 채널 단위로, V는 토큰 단위로 눈금을 따로 잡는 것이 그 때문이다. FP8은 대체로 품질 손실이 작지만, 4비트는 긴 문맥 작업에서 자기 평가 묶음으로 전후를 재 보고 켠다.

PagedAttention과 연속 배칭

KV 캐시를 줄여도 그 메모리를 알뜰하게 쓰지 못하면 이득이 사라진다. 요청마다 최대 길이만큼 연속된 메모리를 미리 잡아 두면 실제로는 절반도 안 쓰는 자리가 생긴다. vLLM이 도입한 PagedAttention은 KV 캐시를 운영체제의 페이지처럼 작은 블록으로 나눠 필요할 때마다 붙인다. 논문은 기존 시스템이 KV 메모리의 60~80%를 이런 식으로 낭비하고 있었다고 보고했다.

여기에 연속 배칭이 겹친다. 배치 전체가 끝날 때까지 기다리지 않고, 한 요청이 끝나는 그 걸음에 대기 중인 새 요청을 빈자리에 넣는 방식이다. 메모리 절감이 처리량으로 바뀌는 경로는 이렇게 이어진다 — GQA와 양자화가 토큰당 바이트를 줄이고, PagedAttention이 그 바이트를 낭비 없이 채우고, 연속 배칭이 빈자리를 곧바로 새 요청으로 메운다. 앞 절에서 본 것처럼 디코딩은 대역폭이 정하므로, 한 걸음에 더 많은 시퀀스를 태울수록 가중치 읽기 한 번이 더 많은 토큰으로 나뉘어 초당 토큰 수가 오른다.

셋 중 하나만 빠져도 효과가 크게 준다. KV를 줄였는데 메모리를 최대 길이로 미리 잡으면 남는 공간이 쓰이지 않고, 블록 단위로 잘 채웠는데 배치를 통째로 기다리면 짧은 요청이 긴 요청을 기다리며 자리를 비워 둔다. 모델 구조의 선택과 서빙 엔진의 선택이 한 줄로 이어져 있는 셈이다. 다음 글에서 다룰 MoE는 방향을 바꿔, 파라미터 수와 토큰당 연산을 떼어 놓는 방식으로 모델을 키운다.

---

읽어주셔서 감사합니다. 😊