MQA와 GQA: 여러 Query가 KV를 공유하기

Q·K·V projection과 Core Attention에서는 각 헤드가 자신의 Query, Key, Value로 정보를 모으는 과정을 살펴봤습니다. 이 구성에서는 쿼리 헤드가 늘어나면 그에 대응하는 키 헤드와 밸류 헤드도 함께 늘어납니다.

토큰을 하나씩 생성할 때는 앞선 위치의 Key와 Value를 저장해 재사용할 수 있습니다. 그런데 문맥이 길어지면 저장할 값도 많아집니다. 여러 Query가 같은 Key와 Value를 함께 사용한다면, Query별로 정보를 모으는 계산을 유지하면서 저장량을 줄일 수 있지 않을까요?

이번 글에서는 먼저 세 구조의 공유 범위를 비교하겠습니다. 이어 MHA에서 토큰 하나가 여러 헤드의 계산으로 이어지는 과정을 짚고, KV를 모두 공유하는 MQA와 그룹 안에서 공유하는 GQA로 확장하겠습니다. 마지막에는 정확히 무엇이 얼마나 줄어드는지와 그에 따른 절충을 살펴보겠습니다.

Query와 KV의 연결 관계

MHA(Multi-Head Attention)는 쿼리 헤드마다 별도의 키 헤드와 밸류 헤드를 사용합니다. MQA(Multi-Query Attention)는 여러 쿼리 헤드가 하나의 키 헤드와 밸류 헤드를 공유합니다. GQA(Grouped-Query Attention)는 쿼리 헤드를 그룹으로 나누고, 같은 그룹 안에서 Key와 Value를 공유합니다.

앞으로 Key와 Value를 함께 가리킬 때는 KV, 대응하는 키 헤드와 밸류 헤드 한 쌍을 가리킬 때는 KV 헤드 하나라고 부르겠습니다. 다음 그림은 쿼리 헤드가 네 개일 때 KV 헤드를 네 개, 두 개, 한 개 사용하는 경우를 비교합니다.

MHA, GQA, MQA를 나란히 비교한다. 세 방식 모두 현재 위치 p3의 q1, q2, q3, q4가 있다. MHA는 q1부터 q4가 각각 K1/V1부터 K4/V4를 읽는다. GQA는 q1과 q2가 KA/VA를, q3과 q4가 KB/VB를 읽는다. MQA는 네 Query 모두 하나의 K/V를 읽는다. 각 KV 묶음은 p0부터 p3까지의 캐시를 뜻한다.

그림의 $q_{1}$ 부터 $q_{4}$ 는 현재 토큰의 4개 쿼리 헤드에서 계산한 Query 벡터입니다. 네 토큰의 Query를 하나씩 나열한 것이 아닙니다. 아래쪽 KV 묶음은 각 Query가 참조할 Key와 Value를 나타냅니다. 화살표는 Query가 어느 KV 묶음을 사용하는지 보여줍니다.

MHA에서는 $q_{1}$ 이 $K_{1} / V_{1}$ 을, $q_{2}$ 가 $K_{2} / V_{2}$ 를 사용하는 식으로 일대일 대응합니다. 가운데 GQA에서는 $q_{1}$ 과 $q_{2}$ 가 $K_{A} / V_{A}$ 를 공유하고, $q_{3}$ 과 $q_{4}$ 가 $K_{B} / V_{B}$ 를 공유합니다. MQA에서는 네 Query가 모두 같은 $K / V$ 를 사용합니다. 세 방식 모두 쿼리 헤드는 네 개이고, 달라지는 것은 KV 헤드의 수와 공유 범위입니다. 이 비교는 GQA 논문의 구조도를 네 쿼리 헤드의 예시로 다시 그린 것입니다.

MHA: 헤드마다 별도의 KV 사용하기

한 층에서 현재 위치 $p_{3}$ 를 처리하는 상황을 보겠습니다. $p_{0}$ , $p_{1}$ , $p_{2}$ 는 앞선 위치이고, 현재 위치까지 포함해 참조할 토큰 위치는 네 개입니다. 여기서 $p$ 뒤의 숫자는 토큰의 종류나 헤드 번호가 아니라 문맥 안의 위치를 뜻합니다.

작은 예시를 위해 모델 차원은 8, 쿼리 헤드는 4개, 각 헤드의 차원은 2로 두겠습니다. K와 V도 헤드마다 두 성분을 사용합니다. 이후의 그림에서도 이 조건을 유지합니다. 실제 모델에서 측정한 수치가 아닌 구조 설명용 예시입니다.

위에는 토큰 위치 p0부터 p3까지 각 입력 벡터를 하나의 블록으로 나란히 표시한다. 현재 위치 p3의 전체 벡터를 WQ, WK, WV에 각각 곱한다. 아래에는 head 1부터 4까지 행을 나누어 p3의 q1부터 q4, k1부터 k4, v1부터 v4를 각각 2성분으로 표시한다. 현재 Key와 Value는 각 head의 캐시에 추가되며, 각 Query는 자기 head의 p0부터 p3까지 KV를 읽는다. Head 1의 계산과 네 출력의 이어 붙이기를 보여 준다.

그림 1의 첫 패널에서 $p_{0}$ 부터 $p_{3}$ 까지의 블록은 각각 해당 위치의 입력 벡터 하나입니다. 이 층에 들어오는 벡터를 뜻하며, 블록 하나가 벡터의 성분 하나인 것은 아닙니다. 현재 처리하는 $p_{3}$ 의 8성분 벡터 전체에 서로 다른 가중치 행렬 $W_{Q}$ , $W_{K}$ , $W_{V}$ 를 곱해 Q, K, V를 만듭니다.

MHA 예시에서는 세 투영의 출력이 각각 8성분입니다. 이를 두 성분씩 네 헤드로 구분하면 $q_{1} \dots q_{4}$ , $k_{1} \dots k_{4}$ , $v_{1} \dots v_{4}$ 가 됩니다. 아래쪽의 작은 칸 하나는 벡터 성분 하나이고, 각 행은 하나의 헤드입니다. 예를 들어 $q_{1}$ , $k_{1}$ , $v_{1}$ 은 모두 현재 위치의 헤드 1에 사용할 벡터입니다. 입력 벡터를 먼저 네 조각으로 잘라 서로 다른 헤드에 보내는 것이 아니라, 각 헤드가 입력 전체를 서로 다른 가중치로 투영합니다.

현재 위치를 계산하기 전에 $p_{0}$ 부터 $p_{2}$ 까지의 K와 V를 보관해 두었다고 하겠습니다. 이렇게 이미 처리한 위치의 K와 V를 저장한 것이 KV 캐시입니다. 새로 만든 $k_{1}$ 과 $v_{1}$ 을 헤드 1의 캐시에 추가하면, $K_{1}$ 과 $V_{1}$ 에는 $p_{0}$ 부터 $p_{3}$ 까지 네 위치의 값이 들어 있습니다. 다른 헤드도 같은 방식으로 현재 위치의 값을 추가합니다.

여기서는 현재 위치의 벡터를 소문자 $k_{1}$ , $v_{1}$ 로, 여러 위치를 모은 행렬을 대문자 $K_{1}$ , $V_{1}$ 로 표시합니다. 각 행렬은 위치 네 개 × 헤드당 두 성분, 즉 $4 \times 2$ 입니다. 행을 하나 내려가는 것은 다음 토큰 위치로 이동하는 것이고, $K_{1}$ 에서 $K_{2}$ 로 바꾸는 것은 다른 헤드의 표현을 사용하는 것입니다.

세 번째 패널은 헤드 1의 계산입니다. 두 성분의 $q_{1}$ 을 $K_{1}$ 의 네 행과 각각 내적하면 위치별 점수 네 개가 나옵니다. Core Attention에서 설명한 스케일링과 softmax를 적용하면 합이 1인 가중치 네 개가 됩니다. 현재 $p_{3}$ 는 $p_{0}$ 부터 자기 위치까지 모두 참조할 수 있습니다.

이 가중치로 $V_{1}$ 의 네 행을 가중합하면 두 성분의 출력 $o_{1}$ 을 얻습니다. 여기서 가중치는 이번 Query가 각 위치의 Value를 얼마나 반영할지 정한 비율입니다. Q·K·V를 만드는 학습 파라미터 $W_{Q}$ , $W_{K}$ , $W_{V}$ 와는 구별해야 합니다.

나머지 세 헤드도 자신의 Query와 KV로 $o_{2}$ , $o_{3}$ , $o_{4}$ 를 계산합니다. 네 출력을 이어 붙이면 $2 + 2 + 2 + 2 = 8$ 성분이 되고, 출력 투영 $W_{O}$ 가 이 성분들을 조합합니다. 그림의 $a_{1}$ , $b_{1}$ 등은 각 출력 벡터의 두 성분을 나타냅니다.

다음 위치를 계산할 때도 과거의 K와 V는 다시 사용하지만, 과거 Query는 이 KV 캐시에 저장할 필요가 없습니다. 새 위치의 Query로 저장된 K와 V를 읽기 때문입니다. 이제 이 계산에서 저장하는 KV 묶음의 수를 줄여 보겠습니다.

MQA: 모든 Query가 KV 공유하기

MHA에서 각 헤드는 위치마다 서로 다른 K와 V를 만들었습니다. MQA는 위치마다 K 벡터 하나와 V 벡터 하나를 만들고, 이를 모든 쿼리 헤드가 공유합니다. Query는 계속 헤드별로 만듭니다.

예시의 현재 위치 $p_{3}$ 에서도 Q 투영은 두 성분씩 네 Query를 만듭니다. 반면 K 투영과 V 투영은 각각 두 성분만 만듭니다. 모델 차원이 8이라면 $W_{Q}$ 의 크기는 $8 \times 8$ , $W_{K}$ 와 $W_{V}$ 의 크기는 각각 $8 \times 2$ 가 되는 구성입니다.

네 Query가 하나의 K 4×2와 V 4×2를 공유한다. 각 Query는 별도의 Attention을 거쳐 o1부터 o4까지 네 출력을 만든다.

그림 2의 왼쪽에는 $4 \times 2$ 크기의 K와 V가 하나씩 있습니다. KV 헤드가 하나라는 것은 저장한 토큰 위치가 하나라는 뜻이 아닙니다. 공유하는 K와 V에도 $p_{0}$ 부터 $p_{3}$ 까지 네 위치의 값이 모두 남아 있습니다.

오른쪽의 $q_{1}$ 은 이 K의 네 행으로 점수를 계산하고, softmax를 거쳐 V를 가중합해 $o_{1}$ 을 만듭니다. $q_{2}$ 도 같은 K와 V를 사용하지만 Query 벡터가 다르므로 위치별 점수와 반영 비율은 달라질 수 있습니다. 예를 들어 한 Query는 $p_{0}$ 의 Value를 더 많이 반영하고, 다른 Query는 $p_{2}$ 의 Value를 더 많이 반영할 수 있습니다.

따라서 같은 KV를 공유한다고 출력까지 같아지는 것은 아닙니다. 네 Query는 각각의 Attention 계산을 거쳐 네 출력을 만듭니다. 그림의 네 Attention 상자는 이 계산 관계를 구별한 것이며, 반드시 별도의 커널을 네 번 실행한다는 뜻은 아닙니다. 출력은 MHA와 마찬가지로 이어 붙인 뒤 출력 투영으로 전달합니다.

저장량은 크게 달라집니다. MHA는 K와 V 묶음을 네 개 저장했지만 MQA는 하나만 저장합니다. 이 예시에서는 $K 4 \times 2 + V 4 \times 2 = 16$ 성분입니다. MQA는 이처럼 헤드 사이에서 KV를 공유해, 토큰을 하나씩 생성할 때 저장하고 읽어야 할 KV의 양을 줄이는 구조입니다. MQA 원논문

GQA: 그룹 안에서 KV 공유하기

GQA는 KV를 공유할 범위를 그룹별로 정합니다. 네 Query를 두 그룹으로 나눈다면, 첫 그룹은 $q_{1}$ 과 $q_{2}$ , 두 번째 그룹은 $q_{3}$ 과 $q_{4}$ 입니다. 각 그룹은 별도의 K와 V를 사용합니다.

q1과 q2는 KA/VA를, q3과 q4는 KB/VB를 공유한다. 각 K와 V는 4×2다. 각 Query는 별도의 Attention을 거쳐 o1부터 o4까지 네 출력을 만든다.

그림 3에서 $q_{1}$ 과 $q_{2}$ 는 $K_{A} / V_{A}$ 를 함께 읽고, $q_{3}$ 과 $q_{4}$ 는 $K_{B} / V_{B}$ 를 함께 읽습니다. 각 그룹의 KV에는 네 위치가 모두 들어 있습니다. 그룹을 나누는 기준은 쿼리 헤드이며, 토큰 위치를 두 그룹에 나누어 저장하는 것이 아닙니다.

현재 위치에서 Q 투영이 만드는 것은 여전히 두 성분씩 네 Query입니다. K와 V 투영은 각각 두 성분씩 두 헤드를 만듭니다. 이 예시의 $W_{K}$ 와 $W_{V}$ 는 각각 $8 \times 4$ 이며, 각 KV 헤드에서 새로 만든 값을 해당 캐시에 추가합니다.

한 그룹 안에서는 MQA와 같은 공유 관계가 성립합니다. $q_{1}$ 과 $q_{2}$ 는 같은 KV를 읽지만 각자 계산한 비율로 Value를 모읍니다. 두 번째 그룹도 동일한 과정을 수행합니다. 그래서 출력은 계속 $o_{1}$ 부터 $o_{4}$ 까지 네 개이고, 저장하는 KV는 두 묶음입니다. 전체 저장량은 $2 \times 16 = 32$ 성분입니다.

쿼리 헤드 수를 $h_{q}$ , KV 헤드 수를 $h_{kv}$ 라고 쓰겠습니다. 같은 크기의 그룹을 만드는 일반적인 구성에서는 한 KV 헤드를 $h_{q} / h_{kv}$ 개의 쿼리 헤드가 공유합니다. 이 예시는 $4 / 2 = 2$ 개입니다. $h_{kv} = 1$ 이면 MQA, $h_{kv} = h_{q}$ 이면 MHA와 같은 공유 구조가 됩니다. 이 두 끝점 사이에서 공유 정도를 조절하는 것이 GQA입니다. GQA 논문 §2.2

KV 헤드 수와 저장량

이제 세 방식의 KV 캐시를 같은 크기의 칸으로 비교하겠습니다. 각 칸은 저장하는 성분 하나이고, 모든 방식에서 토큰 위치는 네 개, K와 V의 헤드당 차원은 각각 2입니다.

T=4에서 MHA, GQA, MQA의 KV 캐시를 동일 크기 칸으로 비교한다. 각 KV head에는 p0부터 p3까지 네 행과 K 두 성분, V 두 성분이 있다. MHA는 4개 head 64성분, GQA는 2개 head 32성분, MQA는 1개 head 16성분이다.

MHA에서는 네 KV 헤드 각각에 K 여덟 성분과 V 여덟 성분이 있어 총 64성분입니다. GQA는 두 KV 헤드이므로 32성분, MQA는 한 KV 헤드이므로 16성분입니다. 토큰 위치나 헤드 안의 성분 수를 줄인 것이 아니라, 별도로 저장하는 KV 헤드 수를 줄인 결과입니다.

이를 한 요청의 한 층에 대해 일반화하면 다음과 같습니다. $T$ 는 캐시에 들어 있는 토큰 위치 수, $h_{kv}$ 는 KV 헤드 수, $d_{h}$ 는 헤드당 차원입니다. K와 V의 헤드 차원이 같다고 가정합니다.

KV 저장 성분 수 = 2 \times T \times h_{kv} \times d_{h}

앞의 2는 K와 V 두 종류를 함께 저장하기 때문에 붙습니다. 그림의 MHA에 대입하면 $2 \times 4 \times 4 \times 2 = 64$ 이고, GQA와 MQA는 $h_{kv}$ 만 각각 2와 1로 바꿉니다. 나머지 조건과 성분당 저장 크기가 같다면, MHA 대비 GQA의 KV 저장량 비율은 $h_{kv} / h_{q}$ 입니다. 이 예시에서 GQA는 절반, MQA는 4분의 1입니다.

이 식은 KV 값 자체의 양을 계산한 것입니다. 모델 전체로 확장하려면 층별 저장량을 더해야 하며, 여러 요청을 처리한다면 요청별 저장량도 더해야 합니다. 실제 메모리 할당에는 구현에 따른 추가 공간이 있을 수 있습니다.

공유 후에도 각 Query는 참조할 위치들의 Key와 점수를 계산하고 Value를 가중합합니다. 쿼리 헤드 수와 문맥 길이를 유지하면 이 Core Attention의 기본 연산량은 같은 규모로 남습니다. 반면 KV 투영의 출력 크기와 저장할 KV의 양은 줄어듭니다. 공유 KV를 잘 재사용하는 구현에서는 메모리에서 가져오는 데이터의 양도 줄일 수 있습니다.

그러므로 캐시가 4분의 1이라고 전체 실행 시간이 반드시 4분의 1이 되는 것은 아닙니다. 실행 시간에는 다른 연산, 모델 가중치 읽기, 커널의 KV 재사용 방식도 영향을 줍니다. 이 글의 64·32·16은 저장 성분 수를 비교한 값입니다.

공유 범위를 정할 때의 절충

MHA에서는 각 쿼리 헤드에 맞는 K와 V 표현을 별도로 학습할 수 있습니다. MQA에서는 모든 Query가 같은 K와 V 표현을 사용해야 하고, GQA에서는 그 제약이 그룹 안에 적용됩니다. KV를 더 넓게 공유할수록 저장량은 줄지만, 헤드마다 서로 다른 K와 V를 만드는 자유도도 줄어듭니다.

이 때문에 MHA로 학습된 모델의 KV를 실행 중 임의로 합친다고 같은 출력을 보존할 수는 없습니다. MQA와 GQA는 KV 투영과 헤드의 연결 관계가 다른 모델 구조입니다. GQA 논문은 기존 MHA 체크포인트의 KV 투영을 변환한 뒤 추가 학습하는 방법을 제시했으며, 해당 실험에서 GQA가 MHA에 가까운 품질과 MQA에 가까운 속도를 얻었다고 보고합니다. 이는 논문의 모델과 학습·실행 조건에서 얻은 결과입니다. GQA 논문 §2–3

따라서 실제 모델의 구성을 볼 때는 쿼리 헤드 수와 KV 헤드 수를 함께 확인해야 합니다. 두 수의 비율로 공유 범위와 KV 저장량을 파악할 수 있고, 품질과 속도는 그 구조로 학습한 모델과 실행 환경에서 확인해야 합니다.

이번 글에서는 Query별 계산과 출력은 유지하면서 KV 헤드의 공유 범위를 바꾸는 방법을 살펴봤습니다. 다음 글에서는 공유 범위에서 한 걸음 더 나아가, 저장하는 KV 표현 자체를 더 작게 만드는 방법을 살펴보겠습니다.