질의 헤드는 그대로, 키·값 헤드만 줄인다

한 줄 요약

MHA·MQA·GQA 를 가르는 손잡이는 하나뿐이다 — 키·값 헤드를 몇 개 둘 것인가. 질의 헤드 수는 건드리지 않는다. 그래서 줄어드는 것은 KV 캐시가 차지하는 메모리이고, 어텐션의 곱셈 횟수는 거의 그대로다.

왜 이게 필요했나

생성은 한 토큰씩 나아간다. 새 토큰을 하나 만들 때마다 모델은 지금까지 나온 모든 자리의 키와 값을 다시 본다. 그래서 그 키와 값을 지워 두었다가 다시 만들지 않고 들고 있는다. 그게 KV 캐시다.

문제는 이 캐시가 문맥이 길어질수록, 층이 많을수록, 동시에 처리하는 요청이 많을수록 그대로 커진다는 점이다. 모델 가중치는 문맥 길이와 무관하게 고정인데 캐시는 아니다. 가중치를 다 올리고 남은 자리를 캐시가 먹으므로, 어느 지점부터는 캐시가 동시 처리 수를 정한다.

게다가 한 토큰을 만들 때 하는 계산은 얼마 안 된다. 그 적은 계산을 하려고 캐시 전체를 메모리에서 읽어 와야 한다. [MQA 원논문](https://arxiv.org/abs/1911.02150)이 증분 디코딩의 병목으로 든 것이 바로 이 되읽기다. 제목이 그대로 답이다 — 쓰는 헤드는 하나면 된다.

여기서 MQA 의 결정이 나온다. 키·값 헤드를 하나만 둔다. 질의 헤드는 여덟 개든 예순네 개든 그대로 두고, 그 모두가 같은 K·V 를 본다. 캐시는 질의 헤드 수분의 1 로 줄어든다.

대신 잃는 것이 있다. 헤드마다 다른 곳을 보라고 만들어 둔 구조에서 "무엇을 내놓을지" 를 정하는 쪽을 하나로 묶어 버린 것이다. [GQA 원논문](https://arxiv.org/abs/2305.13245)은 그 사이에 자리를 만든다. 질의 헤드를 몇 개의 무리로 나누고, 무리마다 키·값 헤드를 하나씩 둔다.

세 이름은 한 축 위의 세 자리

질의 헤드가 h개, 무리가 g개라고 하자. 이때 키·값 헤드 수가 곧 g다.

| 무리 수 g | 이름 | 키·값 헤드 | 한 K·V 를 보는 질의 헤드 |
| --- | --- | --- | --- |
| g = h | MHA | h개 | 1개 |
| 1 < g < h | GQA | g개 | h/g 개 |
| g = 1 | MQA | 1개 | h개 |

세 줄을 따로 외울 것이 아니다. 하나의 손잡이를 어디까지 돌렸는가의 차이다. 조건은 하나뿐이다 — h 가 g 로 나누어떨어져야 한다. 안 그러면 무리 크기가 들쭉날쭉해지고, 어느 무리는 세 헤드를 어느 무리는 두 헤드를 묶게 된다. 실제 설정 파일에서 이 조건을 어기면 대개 모델을 올리는 자리에서 막힌다.

무엇이 줄고 무엇이 안 주나

여기가 이 주제의 핵심이고, 가장 많이 틀리는 자리다. 식을 세워 보면 바로 보인다.

캐시에 들어가는 실수의 개수는 이렇다.

KV 캐시 원소 수 = 2 × 층 수 × (키·값 헤드 수) × 길이 × 헤드 차원

앞의 2 는 K 와 V 두 벌이다. 질의 헤드 수가 이 식에 없다. 질의는 캐시에 남기지 않기 때문이다. 그러니 키·값 헤드를 여덟에서 하나로 줄이면 캐시는 정확히 여덟 분의 일이 된다.

이제 곱셈 횟수를 세어 보자. 길이 n, 질의 헤드 h, 헤드 차원 d 인 한 층을 한 번 통과할 때다.

점수 Q·Kᵀ    h × n × n × d       ← g 가 없다가중합 × V   h × n × n × d       ← g 가 없다K 투영       n × d_model × (g × d)V 투영       n × d_model × (g × d)

어텐션 본체에 g 가 없다. 이유는 구현을 보면 분명하다. 접어 둔 키·값 헤드를 쓸 때는 질의 헤드 수만큼 되풀이해 펴 놓고 평소대로 돈다. 서로 다른 값의 가짓수가 줄었을 뿐, 어텐션이 도는 횟수는 질의 헤드 수가 정한다. 줄어드는 항은 K·V 투영 둘뿐이고, 그건 전체에서 큰 몫이 아니다.

그래서 이렇게 정리된다. GQA 가 줄이는 것은 메모리다. 빨라지는 것은 그 메모리를 덜 읽게 되어서이지 곱셈을 덜 해서가 아니다. 이 구분을 못 하면 "GQA 로 바꿨는데 왜 프리필이 그대로냐" 는 질문에 답할 수 없다.

무리를 어떻게 합칠 것인가

이미 학습된 멀티헤드 모델을 GQA 로 옮긴다고 하자. 무리 안의 키·값 헤드 네 개를 하나로 만들어야 하는데, 어떻게 만들 것인가?

하나만 남기고 버릴 수도 있고, 평균 낼 수도 있고, 처음부터 줄어든 크기로 다시 학습할 수도 있다. GQA 원논문이 고른 길은 무리 안의 키·값 헤드를 평균 내어 초깃값으로 삼고, 거기서 조금 더 학습하는 것이다. 평균만 내고 끝내지 않는다는 점이 중요하다 — 평균은 출발점이지 답이 아니다.

다음 실습도 평균을 쓴다. 다만 학습은 하지 않는다. 그래서 이 실습에서 나오는 "출력이 이만큼 달라졌다" 는 숫자는 품질이 그만큼 나빠진다는 뜻이 아니라, 같은 가중치를 두고 키·값만 묶으면 답이 달라진다는 사실을 보여 주는 값이다. 이 구분을 보고서에도 적게 된다.

현장에서 만나는 모습

첫째, 설정 파일에 헤드 수가 두 줄로 적혀 있다. 요즘 모델 설정에는 질의 헤드 수와 키·값 헤드 수가 따로 있다. 두 값이 같으면 MHA, 키·값 쪽이 1 이면 MQA, 그 사이면 GQA 다. 세 이름을 찾을 것 없이 두 숫자의 비를 보면 된다.

둘째, "GQA 인데 왜 안 빨라지나" 가 나온다. 프리필처럼 한 번에 많은 자리를 처리할 때는 계산이 병목이고, 거기에는 g 가 거의 영향을 주지 않는다. 이득은 긴 문맥을 들고 토큰을 하나씩 뽑는 구간, 그리고 동시 요청 수를 올리는 데서 나온다.

셋째, 동시 처리 수의 상한이 캐시로 정해진다. 남은 메모리를 캐시 한 벌 크기로 나눈 값이 대략의 상한이다. 키·값 헤드를 넷으로 줄이면 그 한 벌이 작아지므로 상한이 올라간다. 이때 늘어나는 것은 처리량이지 한 요청의 속도가 아니다.

넷째, 텐서 병렬에서 키·값 헤드 수가 걸림돌이 된다. 헤드를 여러 장치에 나눠 실을 때, 키·값 헤드가 장치 수보다 적으면 나눌 것이 모자라 같은 키·값을 여러 장치가 들고 있게 된다. 캐시를 줄이려고 g 를 낮췄는데 장치마다 복제되면 기대만큼 안 준다.

다섯째, 헤드 수를 바꾸면 다른 모델이다. 무리를 평균해 만든 가중치를 그대로 서비스에 올리면 출력이 달라진다. 원논문도 옮긴 뒤 추가 학습을 한다. "설정만 바꾸면 된다" 가 아니다.

실무에서 진짜 중요한 것

• 줄어드는 것은 메모리다. 곱셈 횟수는 거의 그대로다. 이득이 어디서 오는지 알아야 어디에 쓸지 정할 수 있다.
• h 를 g 로 나눌 수 있어야 한다. 무리 크기가 고르지 않으면 구현이 성립하지 않는다.
• 두 끝은 같은 축 위에 있다. g=h 면 MHA, g=1 이면 MQA. 코드 한 벌로 셋을 다 돌릴 수 있게 짜 두면 견주기가 쉬워진다.
• 옮기는 것과 학습하는 것은 다른 일이다. 평균은 초깃값일 뿐이다.

다음 실습에서 할 것

/root/work/tf-gqa/gqa.py 를 한 단계씩 키운다. 질의 헤드를 이어진 덩어리로 무리 짓는 것에서 시작해, 무리 안의 키·값 헤드를 평균으로 접고, 그것을 질의 헤드 수만큼 되풀이해 편 뒤, 한 헤드짜리 어텐션을 거쳐 세 방식을 같은 가중치로 돌린다. g 를 h 로 두면 MHA 와 한 자리도 다르지 않아야 하고, g 를 줄이면 출력이 달라져야 한다. 그 차이를 숫자로 잰다.

그다음이 이 실습의 요점이다. KV 캐시 원소 수와 곱셈 횟수를 각각 세어 나란히 놓는다. 캐시는 g 에 그대로 비례해 줄어드는데 어텐션 본체의 곱셈은 한 번도 줄지 않는다는 것을, 남의 설명이 아니라 자기가 센 숫자로 확인하게 된다. 마지막에 그 두 표와 출력 차이를 보고서로 남긴다.

시간은 재지 않는다. 이 파드에는 GPU 가 없고 CPU 도 다른 작업과 나눠 쓰므로 "빨라졌다" 는 여기서 잴 수 있는 것이 아니다. numpy 도 시스템 파이썬에는 없다(/opt/onnx-lab 안에만 있다). 표준 라이브러리만으로 충분하고, 판정은 전부 원소 수·곱셈 횟수 같은 정수 세기와 허용 오차를 둔 수치 대조다. 채점기는 여러분의 모듈을 실제로 불러 매번 다른 헤드 수와 무리 수로 함수를 두드려 보고, 따로 계산한 값과 대조한다.