딥시크 MLA 완전 해설: KV 캐시 57배 절감과 6배 빠른 생성의 원리

핵심 메시지

• 딥시크가 2024년 6월 발표한 멀티헤드 잠재 어텐션(MLA)은 트랜스포머의 주변부가 아니라 어텐션이라는 핵심을 건드린 개선이다.
• MLA는 KV 캐시 크기를 57분의 1로 줄이면서 딥시크 구현 기준으로 기존 트랜스포머보다 6배 이상 빠르게 텍스트를 생성하게 만들었다.
• KV 캐싱은 어텐션 연산량을 토큰 수의 제곱에서 선형으로 낮춰 주지만, 그 대가로 막대한 메모리 대역폭을 요구한다.
• 멀티쿼리 어텐션과 그룹쿼리 어텐션도 캐시를 줄이지만 헤드가 키·값을 공유하면서 성능 손실을 감수해야 했다.
• MLA의 핵심 아이디어는 키와 값을 압축된 잠재 공간에 저장하고, 헤드별 고유 가중치로 다시 펼치되 그 가중치를 질의와 출력 계산에 흡수시켜 추론 시 추가 연산을 없앤 것이다.

쉽게 이해하기

2025년 1월 중국 기업 딥시크는 선도 모델의 일부 연산량만으로 경쟁 가능한 성능을 낸 R1을 공개해 업계를 놀라게 했다. 더 인상적이었던 것은 공개 범위였다. 모델 가중치와 추론 코드, 그리고 상세한 기술 보고서를 함께 내놨고, 2024년 한 해 동안 평균 월 1편꼴로 보고서를 발표하며 R1으로 이어지는 혁신들을 축적해 왔다. 그중 2024년 6월에 소개된 것이 멀티헤드 잠재 어텐션, 줄여서 MLA다.

MLA를 이해하려면 어텐션이 무엇을 계산하는지부터 봐야 한다. 언어 모델은 토큰을 하나씩 생성하며, 각 토큰은 앞선 모든 토큰의 함수다. 토큰 간 상호작용은 어텐션 패턴이라는 행렬로 처리된다. GPT-2 스몰은 층마다 12개 헤드, 12개 층으로 총 144개의 패턴을 만들고, 딥시크 R1은 층마다 128개 헤드에 61개 층으로 총 7,808개의 패턴을 만든다. 입력 토큰이 9개면 패턴은 9×9 행렬이 된다.

영상은 '미국 국기는 빨강 하양 그리고'라는 예문으로 어텐션 헤드의 역할을 보여준다. 어떤 헤드는 '미국'이라는 수식어를 '국기'에 적용해 하나의 개념으로 묶고, 또 다른 헤드는 국기·빨강·하양을 마지막 토큰 위치로 끌어와 다음 단어 '파랑'을 정확히 예측하도록 돕는다. 계산 과정은 입력 행렬에 학습된 가중치를 곱해 질의(Q)와 키(K)를 만들고, 둘의 내적으로 유사도를 구한 뒤 마스킹과 소프트맥스를 거쳐 패턴을 얻고, 여기에 값(V) 행렬을 곱해 데이터를 가공하는 순서다.

문제는 비용이다. 어텐션 패턴의 크기는 입력 토큰 수의 제곱으로 커진다. 10만 토큰 컨텍스트라면 해리 포터 1권 분량의 텍스트를 한 줄과 한 열로 늘어놓고 모든 토큰 쌍의 내적을 구하는 셈이다. 다행히 지름길이 있다. 토큰을 하나 더 붙여도 앞선 키와 값은 그대로이므로, 이를 메모리에 저장해 두고 새로 생긴 마지막 행만 계산하면 된다. 이것이 KV 캐싱이고, 덕분에 어텐션 연산량은 제곱이 아니라 선형으로 늘어난다.

대신 메모리가 폭발한다. 모든 층, 모든 헤드에 대해 세션 전체의 키와 값을 들고 있어야 하기 때문이다. 영상의 계산에 따르면 딥시크 R1 구조에 10만 토큰 컨텍스트라면 토큰당 4메가바이트, 즉 새 토큰 하나를 만들 때마다 400기가바이트를 읽어야 한다. 기존 해법인 멀티쿼리 어텐션은 모든 헤드가 키·값을 공유해 캐시를 헤드 수만큼 줄이지만 헤드의 전문화를 막아 성능이 떨어졌고, 라마 3가 쓰는 그룹쿼리 어텐션은 8개씩 묶어 8분의 1로 줄였지만 여전히 손실이 있었다.

MLA는 다른 길을 택했다. 입력을 모든 헤드가 공유하는 압축된 잠재 공간으로 먼저 투영하고, 거기서 헤드마다 서로 다른 학습 가중치로 키와 값을 복원하는 방식이다. 언뜻 행렬 곱이 하나 늘어난 것처럼 보이지만, 이 가중치들은 학습이 끝나면 고정되므로 선형대수적으로 질의 계산과 최종 출력 계산에 미리 흡수시킬 수 있다. 결과적으로 추론 시 추가 연산 없이 잠재 캐시에서 곧바로 어텐션 패턴을 구한다. 캐시 크기는 헤드 수와 무관해지고 잠재 차원에만 의존하게 되며, R1 기준 토큰당 70킬로바이트로 줄어든다.

주요 인사이트

• MLA의 성과가 특별한 이유는 절충을 피했다는 데 있다. 멀티쿼리와 그룹쿼리 어텐션은 캐시를 줄이는 대신 성능을 내줬지만, MLA는 캐시를 57분의 1로 줄이면서 오히려 성능을 개선했다.
• 핵심 아이디어는 머신러닝에서 익숙한 잠재 공간을 어텐션 캐시에 적용한 것이다. 무엇을 공유할지 사람이 정하는 대신, 모델이 키와 값을 어떻게 압축하고 헤드 사이에 나눌지 스스로 학습하게 만들었다.
• 가중치 흡수 기법은 실무적으로 중요한 교훈을 준다. 학습 시점에 고정되는 선형 변환은 추론 경로에서 미리 합쳐 둘 수 있고, 그 결과 '메모리를 아끼려다 연산을 늘리는' 흔한 함정을 피할 수 있다.
• 질의는 캐싱할 필요가 없다는 점도 짚어 둘 만하다. 어텐션 패턴을 갱신할 때 필요한 것은 새로 생긴 마지막 질의 행뿐이므로 캐시 대상은 키와 값에 한정된다.
• 숫자로 보면 차이가 분명하다. 같은 R1 구조에서 전통적 어텐션은 토큰당 4메가바이트, 그룹 크기 8의 그룹쿼리 어텐션은 500킬로바이트, MLA는 70킬로바이트를 요구한다.

자주 묻는 질문

관련 AI 소식