AI VIDEO BRIEFING

MQA GQA MLA 완전 정리: KV 캐시 메모리 폭증을 잡는 어텐션 변형 세 가지의 원리와 대가

LLM 추론 속도를 살리려 만든 KV 캐시가 이번엔 GPU 메모리를 잠식한다. MQA, GQA, 그리고 딥시크가 쓴 MLA가 각각 무엇을 줄이고 그 대가로 무엇을 내주었는지, 대학 LLM 강의 내용을 따라 차례로 정리했다.

KV 캐시가 GPU를 잡아먹는다, MQA·GQA·MLA는 무엇을 포기했나 영상 대표 이미지

핵심 메시지

  • 언어모델은 토큰을 하나씩 만들기 때문에 이전 문맥의 키와 값을 매번 다시 계산해야 하고, 이를 저장해 두는 것이 KV 캐시다.
  • KV 캐시는 연산량을 크게 줄이지만 그 대가로 GPU 메모리를 먹으며, 그 크기는 문맥 길이에 비례해 걷잡을 수 없이 커진다.
  • MQA는 키와 값의 헤드를 하나로 줄여 캐시를 헤드 수만큼 줄이고, GQA는 그 사이에 그룹 개수라는 조절 손잡이를 놓는다.
  • 딥시크의 MLA는 키와 값을 그대로 저장하는 대신 훨씬 작은 잠재 코드로 압축해 저장하고, 쓸 때 다시 펴는 방식이다.
  • MLA는 추가 행렬 곱으로 인한 지연을 가중치 흡수로 없애지만, 회전 위치 임베딩과 충돌해 별도 가지를 두는 우회가 필요했다.

쉽게 이해하기

강의는 셀프 어텐션의 계산 구조를 다시 짚는 데서 출발한다. 문장을 한 토큰씩 만들어 나갈 때 쿼리는 지금 처리 중인 토큰 하나에만 해당하지만, 키와 값은 그 앞에 쌓인 문맥 전체에 대해 만들어야 한다. 토큰이 하나 늘 때마다 앞의 모든 토큰에 대해 가중치 행렬을 다시 곱하는 것은 명백한 낭비이고, 여기서 이미 만들어 둔 키와 값을 저장해 두고 재사용하자는 발상이 나온다. 그것이 KV 캐시이며, 이 구조는 추론을 두 단계로 갈라놓는다. 사용자가 던진 질문은 이미 전부 알고 있으므로 한 번에 병렬로 처리해 캐시를 채우는 프리필 단계를 거치고, 그다음 답변을 한 토큰씩 생성하면서 캐시를 읽고 다시 채워 넣는 디코드 단계로 넘어간다.

문제는 트레이드오프다. 강의는 KV 캐시 용량을 정밀도, 키·값 차원, 헤드 수, 레이어 수, 시퀀스 길이의 곱으로 계산해 보인다. 딥시크 V3 설정에 2바이트 정밀도, 키·값 차원 128, 헤드 128개, 레이어 61개, 문맥 3만 2천 토큰을 넣으면 이론상 최대 130GB가 나온다. 디스크가 아니라 GPU 메모리 130GB이며, 그 안에는 모델 가중치도 활성값도 함께 올라가야 한다. 문맥 길이가 백만 토큰 단위로 늘어나는 흐름을 생각하면 이대로 둘 수 없는 숫자다.

첫 번째 해법인 MQA는 2019년에 나왔다. 쿼리만 멀티헤드를 유지하고 키와 값은 헤드를 하나만 쓴 뒤 복제해 쓰자는 것으로, 캐시가 헤드 수만큼 줄어든다. 대신 키와 값이 단일 헤드가 되면서 성능 저하를 감수해야 한다. GQA는 그 사이를 메운다. 헤드를 1로 줄이는 대신 미리 정한 그룹 개수만큼만 만들어, 메모리와 성능 사이의 지점을 하이퍼파라미터로 고르게 한다. 강의에 인용된 비교에서 GQA는 MQA와 비슷한 속도대에서 더 나은 성능을 보였고, 그래서 최근 널리 쓰인다.

MLA는 발상을 한 단계 비튼다. 강의는 먼저 MHA와 GQA를 다운 프로젝션과 업 프로젝션의 조합이라는 공통 형식으로 다시 쓴 뒤, 그렇다면 중간에 놓인 것을 원래 차원 그대로 저장할 이유가 없다고 지적한다. 훨씬 작은 차원의 압축된 잠재 코드를 만들어 그것만 캐싱하고, 어텐션을 돌릴 때 업 프로젝션으로 되살리자는 것이다. 딥시크는 이 잠재 차원을 576으로 잡았는데, 헤드 128개와 키·값 차원 128, 그리고 키와 값 둘을 모두 저장해야 하는 데서 나오는 원래 규모에 비하면 훨씬 작다.

남은 문제는 구현이다. 업 프로젝션 행렬을 추론 때마다 곱하면 지연이 생기는데, 학습이 끝나면 고정되는 프로젝션 행렬들을 미리 하나로 합쳐 두는 가중치 흡수로 이를 없앤다. 그러면 일반적인 멀티헤드 어텐션과 같은 지연으로 캐시 크기만 줄인 셈이 된다. 다만 회전 위치 임베딩은 쿼리와 키를 만든 뒤에 끼어들어 이 합치기를 깨뜨리기 때문에, 위치 정보 전용 가지를 따로 만드는 디커플드 방식으로 우회한다.

주요 인사이트

  • KV 캐시는 연산을 아끼려고 도입한 장치지만, 결과적으로 메모리라는 다른 자원으로 비용을 옮겨놓았을 뿐이다.
  • MHA에서 MQA, GQA, MLA로 이어지는 흐름은 성능을 올리는 경쟁이 아니라 같은 성능을 더 적은 캐시로 내려는 경쟁이다.
  • MLA를 다운 프로젝션과 업 프로젝션이라는 공통 언어로 다시 쓰면, MHA와 GQA도 그 특수한 경우로 정리된다.
  • 이론상의 개념은 단순한데 실제 프로덕션에 넣을 때 필요한 가중치 흡수와 디커플드 위치 임베딩이 구현의 대부분을 차지한다.
  • 강의는 MLA가 KV 캐시를 크게 줄이면서도 MHA에 뒤지지 않는 성능을 보였다고 전하며, 잠재 코드를 거치는 과정 자체가 표현력에 기여했을 가능성을 짚는다.

자주 묻는 질문

KV 캐시는 왜 쿼리는 저장하지 않고 키와 값만 저장하나?

쿼리는 지금 디코딩 중인 토큰 하나에 대응하므로 매 단계 새로 만들 수밖에 없다. 반면 키와 값은 이전 문맥 전체에 대해 필요하고 그 값이 단계마다 바뀌지 않으므로, 한 번 만들어 두면 계속 재사용할 수 있다.

프리필과 디코드는 어떻게 다른가?

프리필은 이미 다 알고 있는 질문 토큰들을 한 번에 병렬 처리해 KV 캐시를 채우는 단계이고, 디코드는 답변을 한 토큰씩 생성하면서 그 캐시를 읽고 새로 만들어진 키와 값을 다시 채워 넣는 단계다.

MLA에서 회전 위치 임베딩이 문제가 되는 이유는?

가중치 흡수는 프로젝션 행렬들을 미리 하나로 합칠 수 있다는 전제 위에 있는데, 회전 위치 임베딩은 쿼리와 키를 만든 뒤에 회전 행렬을 곱하므로 그 사이에 끼어들어 합치기를 막는다. 그래서 위치 임베딩 전용 가지를 따로 두는 디커플드 방식을 쓴다.

원문과 출처

이 글은 원본 영상의 자막을 바탕으로 한국어 독자를 위해 요약했습니다. 전체 맥락과 최신 정보는 원문에서 확인하세요.

YouTube 원본 영상 보기 ↗

관련 AI 소식