기술 발전 순서로 알아보는 KV Cache

안녕하세요, 이번엔 KV Cache에 대해서 공부해본 내용을 포스팅으로 정리해볼까 합니다. 요즘 AI 엔지니어링을 하면서 가장 고려해야하는 것이 바로 이 KV Cache가 점유하게되는 메모리를 관리하는 것인 것 같습니다. 그래서 이번엔 AI 엔지니어링을 하게되면 반드시 듣게되는 KV Cache란 무엇이고 이 KV Cache를 효율적으로 관리하기 위해 탄생한 다양한 기술들에 대해서 정리해보겠습니다. 

 

Attention이란?

KV Cache를 알아보기 전에 먼저 Transformer의 Attention 연산에 대해서 언급하고 넘어가도록 하겠습니다. Transformer의 Attention 매커니즘에서 데이터를 처리할 때 (Inference 할 때) 각 토큰은 세 가지 역할을 부여받습니다. 

 

• Query (Q) : 내가 지금 찾고 있는 정보가 무엇인가? (질문)
• Key (K) : 내 정보는 이런거야 나랑 매칭 해볼래? (인덱스)
• Value (V) : 나랑 매칭됐으면 이 값을 가져가! (값)

이 구조는 우리가 백엔드에서 Map<Key, Value>로 데이터를 가져올 때 get(Key) 이렇게 가져오듯이 모델은 Query를 던져서 과거의 Key들과 대조해보고 연관성이 높은 순서대로 Value를 긁어옵니다. 

 

이 과정을 Attention이라고 부르고 이 기본적인 개념에서부터 KV Cache를 설명해보도록 하겠습니다. 

 

MHA (Multi Head Attention)

기본적인 Transformer의 Attention연산은 MHA를 따르고 이 때 생성된 행렬 값들은 메모리에 저장되고 이를 우리는 KV Cache라고 부릅니다. 

 

여기서 새롭게 등장하는 개념이 바로 헤드인데 여기서 헤드란 Q, K, V에 매칭되는 개념이고 우리는 쿼리헤드, K헤드, V헤드 이런식으로 명칭을 붙입니다. MHA는 가장 원초적인 형태의 Attention 연산을 일컫는 기술이었으며 Q, K, V가 1:1:1로 이루고 있었죠. 

 

쿼리와 Key, Value를 조금 더 직관적으로 이해하자면 다음과 같습니다.

 

• Q : 나한테 명령하는 상사
• K : 실제 일 처리하는 애
• V : 산출물

이렇게 이해하면 한층 더 쉽습니다. 

 

MHA의 개념은 일을 시키는 놈과 실제로 일을 처리하는 애가 1:1로 이루고 있었기 때문에 답변 퀄리티가 굉장히 좋았습니다. 다만 이렇기 때문에 회사에 자리를 많이 차지하는 문제가 있었고 회사가 좁아터지게 되었죠.

 

이 방법은 KV 행렬 값을 메모리에 올려 놓고 토큰을 생성할 때마다 사용했기 때문에 연산을 효율적으로 할 수는 있었으나 메모리를 잡아먹는 속도가 어마어마했습니다. 

 

그래서 이 방법은 금방 잊혀져갔습니다. 

 

MQA (Multi Query Attention) 과 GQA (Grouped Query Attention)

이 방법은 MHA가 KV Cache를 무지막지하게 많이 잡아먹어 고안된 방법들입니다. 

 

MQA

예를 들어서 쿼리 헤드가 8개일 때 (상사가 8명일 때) KV헤드가 단 한개인 구조 (실제 일처리 하는 애는 한명일 때)를 가지고 있는 것을 말합니다. 8개의 서로 다른 쿼리가 하나의 인덱스와 하나의 데이터를 보고 각자가 필요한 것을 찾아가는 방식이죠. 

 

장점은 KV Cache용량이 압도적으로 줄어든다는 것이지만 성능이 압도적으로 낮아진다는 단점이 있습니다. 상사가 8명인데 나한테 각자가 다른 일을 시킨다고 생각하면 정말 끔찍하네요... 

 

일단 KV 헤드가 기억하고 있는 컨텍스트가 부족하고, KV 헤드 하나가 처리할 수 있는 정도도 오버하기 때문에 답변 퀄리티에 상당한 악영향을 미쳤다고 볼 수 있습니다. 

 

GQA

이 방법은 MHA와 MQA의 절충안으로서 쿼리 헤드를 그룹으로 묶어서 그룹 한개당 KV 헤드를 하나씩 할당해주는 방법입니다. 

 

즉, 일을 시키는 상사가 2명에 처리하는 애가 1명인 구조이죠. 

 

이 방법은 MQA에 비하면 메모리는 많이 잡아먹을지언정 생성되는 답변 퀄리티가 괜찮아진다는 장점이 있습니다. 또한, MHA에 비해서는 KV Cache를 덜 먹는다는 것이 또 다른 장점이었죠. 

 

현재 GQA는 메타부터 구글까지 다양한 미국의 빅테크 기업들이 많이 채용하는 KV Cache 관리 전략이 되었습니다. 

 

 

이건 Gemma-3-27b-it의 config.json을 가져온 것입니다. 

 

잘 보시면 num_attention_heads 가 32인 것과 num_key_value_heads 가 16인 것을 확인할 수 있습니다. 

 

쿼리 헤드는 32개에 KV헤드가 16개라는 뜻이고 이는 앞선 예시처럼 2명의 상사와 1명의 부하를 가지고 있다고 볼 수 있죠. 

 

MLA (Multi-Head Latent Attention)

제가 굳이 미국 빅테크라고 언급한 이유가 중국에서는 그래픽카드가 부족해서 온몸비틀기를 할 수 밖에 없었습니다. 메모리 용량이 부족한데 KV Cache가 미친듯이 늘어나는 것을 막아야했거든요. 

 

MLA 방식은 헤드를 여러개 쓰는 것은 동일하지만 저수준 잠재적 벡터 (Intent Vector)로 벡터를 압축해 차원을 줄인 뒤 필요할 때 다시 업프로젝션 (Up-Projection)해서 사용하는 방법입니다. 

 

수학 얘기가 조금 나와서... 사실 이 부분을 포스팅하는게 맞나 싶긴 한데... 

 

우리가 구하려는 Attention값은 사실 머신러닝에서도 많이 등장하는 Score입니다. 이걸 Attention Score라고 부르죠. 

 

이 Attention Score를 구하는 방법은 Query와 Key를 내적하는 것입니다. 

 

만약 고등학교 떄 기하와 벡터를 공부해보신 분들이라면 등장하는 익숙한 수식인

바로 이것입니다. 

 

여기서 중요한건 코사인 값입니다. 

 

두 값을 내적했는데 코사인 값이 1이어서 (세타가 0이어서) 내적의 값이 크게 나왔다면? 이 두 벡터는 동일한 방향을 가리키고 있는 것이라는 뜻입니다. 컴퓨터 공학적으로 이를 해석하면 유사도가 매우 높다라고 판단할 수 있는 것이죠. 

 

코사인 값이 0에 가까워질 수록 (세타가 90도에 가까워질수록) 내적 값이 작아져서 0에 수렴하고 그럼 이 두 벡터는 수직방향을 가리키고 있는 것이라고 판단하고 이를 유사도가 없다. 라고 판단합니다. 

 

자... 이제 본격적으로...

 

Attention Score를 구하기 위해서는 쿼리와 Key를 내적하는 것이라고 그랬죠? 

 

MLA에선 이 내적에 약을 칩니다. 벡터의 차원을 낮추는 약을 치는거죠. 이를 "다운프로젝션을 행렬곱했다." 라고 표현합니다. 그리고 다시 필요할 때 업프로젝션을 행렬곱하면서 원래대로 원상복구 하는 것이죠. 

 

이렇게 업프로젝션을 하면 오버헤드가 발생하지 않을까 살짝 걱정이 되긴 합니다. 하지만 오버헤드는 발생하지 않습니다. 

 

오잉? 어째서죠? 

 

이 벡터를 원래대로 돌리는 복원된 Key의 수식은 다음과 같습니다. 

 

다른건 다 보지 않으셔도 됩니다. W라고 써있는게 업프로젝션 행렬이고 이를 곱한 결과값인 k가 Key 입니다. 

 

그리고 아까 Attention Score는 쿼리와 Key의 내적이라고 했으니 수식은 다음과 같습니다. 

다른건 다 보지 않으셔도 됩니다!!!

 

원래 쿼리에 업프로젝션한 Key를 곱했다~ 라고 생각하시면 됩니다. 

 

이걸 직관적으로 이해하기 위해서는 다음과 같은 예시를 생각해볼 수 있는데요. 

 

만약 100만명이 마라톤을 출전했다고 가정해봅시다. 이 마라토너들의 가슴팍과 등팍에는 숫자가 써있습니다. 자신이 참여한 숫자이죠. 

 

이 숫자를 한번에 다운사이징을 하는거죠! (이유는 묻지 말아주세요)

 

그래서 마라토너들의 몸에 붙어있는 숫자가 보이지 않게 되었습니다. 

 

이 때 결승전에 도착한 선수들의 몸에 붙어있는 숫자로 순위를 매겨야하는데 도착지점에 있는 진행위원이 글자가 너무 작아 보이지 않는다고 하소연을 하기 시작합니다. 

 

지금 저 Attention Score의 수식은 100만명의 모든 출전자에게 "다시 숫자 크기를 키워오세요!!" 라고 말하는 것과 같습니다. 

 

그럼 정말 어마어마한 오버헤드가 생길겁니다. 

 

그래서 딥시크는 행렬곱의 결합법칙을 이용해서 이 문제를 해결했습니다. 결합법칙은 다들 아시다시피 곱셈의 순서를 바꿔도 값이 똑같다는 점을 이용하는 것이죠. 

이렇게 바꿨습니다. 

 

이게 무슨 의미냐하면 100만명의 모든 출전자에게 다시 숫자 크기를 키워오라는게 아니고 진행요원한테 현미경을 하나 쥐어준것과 같습니다. 다시 벡엔드적인 해석을 하자면, 100만개의 행을 풀스캔해서 찾는 것 보다 앞에서 인덱싱해서 찾는 것이 더 빠른 것과 같은 이치입니다. 

 

이번엔 다시 한글로 풀어서..

스코어 = 쿼리 X 업 프로젝션 행렬 X 키

쿼리에 키를 내적한 것이 Attention Score인데 다운프로젝션된 행렬값과 키를 곱한 것을 KV Cache에 저장해놓고 있었는데, 이를 다시 업프로젝션 하기 위해 업프로젝션된 행렬로 복원한 뒤 쿼리와 일일히 비교했어야 했습니다. 

 

원래였으면 업프로젝션 행렬과 키를 먼저 곱하고 쿼리를 곱해야했지만, 이를 결합법칙으로 쿼리와 업프로젝션 행렬을 먼저 곱해버리는 것입니다. 이렇게 되면 업프로젝션 된 쿼리가 튀어나오고 이걸 키를 곱하면 되기 때문에 연산량이 0에 수렴하게 되는 것이죠. 

 

0은 아니고 0에 수렴하는 것은 지금의 예시에서 1/100만이 되었기 때문에 엄연히 0은 아니라는 말입니다. 

 

이걸 행렬 흡수라는 이름으로 딥시크가 적용한 기술이죠.

 

RoPE라는 변수

최근 LLM들은 토큰의 위치를 알기 위해 RoPE (Rotary Positional Embedding)이라는 기술을 사용합니다. 그런데 RoPE는 벡터를 회전시켜야하기 때문에 행렬 흡수 수식과 수학적으로 섞일 수 없었습니다. 

 

왜냐하면 행렬은 회전 후 곱하는 것과 곱한 후 회전하는 것의 결과가 다르기 때문인데요. 

 

그래서 딥시크는 쿼리의 키를 두 부분으로 쪼개서 계산합니다. 

 

첫 번째 항은 압축된 캐시를 이용한 효율적인 문맥을 검색하고, 두 번째 항은 토큰의 위치를 찾기 위해 별도로 관리되는 정보 검색용이라고 생각하시면 됩니다. 이렇게 두개를 더한 값으로 Attention Score를 구하게 되는것이죠. 

 

DSA (DeepSeek Sparce Attention)

최근 DeepSeek V3에서는 DSA라는 기술이 추가적으로 더 공개되었습니다. 

 

기본적인 Attention 연산은 문장이 길어질수록 연산량이 제곱으로 늘어나기 때문에 문맥이 만약 128k만 되더라도 MLA로 감당할 수 없는 문제가 생깁니다. GPU가 부족한 중국은 KV Cache를 방망이 깎는 노인처럼 깎기 시작했습니다. 

 

1. Lightning Indexer : 먼저 모든 토큰을 다 대조하는게 아니라 FP8 정도의 정밀도를 가진 아주 가벼온 인덱서를 먼저 돌립니다. 
2. Top-K Selection : 인덱서가 이 쿼리에는 이 100개의 과거 토큰 (Key-Value)가 가장 중요하다고 Top-K를 찍어주면 모델은 딱 그 부분만 정밀하게 계산합니다. 
3. 그 결과 연산 복잡도가 제곱에서 선형적으로 늘어나게 되었습니다. 

DSA와 MLA는 서로 대체되는 기술이 아니라 MLA구조 위에 구현된 상위 레이어입니다. 

 

MLA는 데이터를 저차원으로 압축해서 메모리를 아끼는 기술이고 DSA는 압축된 데이터 중에서 현재 질문과 가장 관련있는 놈들만 골라서 연산량을 아끼겠다는 취지인 것이죠. 

 

딥시크의 최신 모델들 (V3.2)은 여기서 한발 더 나아가서 MQA와 DSA를 합치려는 움직임도 보입니다. KV 헤드를 공유하는 MQA와 선택전 연산인 DSA를 합쳐서 메모리 사용량을 폭발적으로 아끼려는 움직임이죠. 

 

딥시크는 MLA와 DSA를 적용하면서 문장이 아무리 길어져도 첫 번째 토큰이 나오는 속도인 TTFT를 거의 일정하게 유지할 수 있게 되었습니다. 

 

어른들의 사정..?

왜 이런 좋아보이는 기술이 중국쪽에서 우수수 쏟아지는걸까요? 중국이 후발주자라 혁신적이지 않으면 살아남지 못해서? 중국인이 똑똑해서? 

 

역시 앞에서도 언급했듯이 부족한 그래픽카드가 한 몫을 했을 것입니다. 그래픽카드 서버가 부족해서 나오는 온몸비틀기에 가까운 것이죠. 

 

미국은 중국에 비해 상대적으로 풍부한 GPU서버를 공급받을 수 있었고 중국은 대중 무역 제재로 인해 그렇지 못했습니다. 이 차이가 기술적인 차이를 만든 것이죠. 

 

LLM 모델이 메모리에 올라가는 것보다 KV Cache가 점유하는 메모리가 훨씬 더 큰 (사실 말도 못할정도로 많은 차이가 납니다) 상황에서 KV Cache를 줄이는 것이 곧 필요한 서버의 수가 줄어드는 것이되고 이것이 곧 중국이 살아남기위한 방법이 되는 것입니다. 

 

기존 GQA가 표준이 되었고 많은 미국의 빅테크들이 이 방법을 채택하는 이유는 LLama나 GPT 시리즈에서 수조개의 파라미터를 학습시키고 안정성까지 검증받았기 때문에 굳이 MLA로 넘어갈 필요가 없었던 것입니다. 

 

또한, MLA는 벡터를 압축하고 다시 원상복구하는 복잡한 로직이 필요했고 이 과정은 커널을 뜯어고쳐야하는 수준입니다. 중국에서도 CUDA커널 말고 새로운 커널 시스템이 필요했고 딥시크에선 Triton이라는 이름으로 새로운 커널을 만들면서까지 개발하게 되었습니다. 

 

조금 다른 얘기긴 하지만 며칠전 등장한 카카오의 Kanana-2가 딥시크 아키텍처를 적용했다고 하는데 이는 커널 엔지니어링이 가능한 개발 수준을 갖추었다고 해석해도 되고 한국의 AI 개발 실력이 빅테크를 중심으로 멈춰있지 않다는 것을 보여주는 사례로 볼 수 있습니다. 한국 역시 그래픽카드가 부족했고 그렇기에 이런 기술적인 의사결정을 하지 않았나 싶습니다. 

 

미국 빅테크 입장에선 커널 엔지니어링까지 해야하고 기존에 학습시킨 데이터가 뭉개질 수도 있는 리스크있는 일을 굳이 하지 않은 것이겠죠. 그것보다 그냥 그래픽카드 서버 한대를 더 박아넣는게 경제적으로 더 이득이었기 때문에 이런 선택을 한 것일 수도 있습니다. 

 

때문에, 딥시크에서 개발한 MoE는 저수준 개발이 필요 없었기 때문에 대부분의 빅테크들이 이를 적용한 반면, MLA같은 저수준 개발이 필요한 것들은 선뜻 적용하지 않은 것으로 보입니다. 아마 많은 리스크가 있다고 판단한 것이겠죠. 

 

하지만, 미국 빅테크 입장에서 이를 적용하지 않을 이유는 전혀 없습니다. 10만대로 처리하던 것을 1만대로 처리할 수 있다면야 유지비용도 줄이고 서버 공간과 같은 물리적인 비용까지 줄일 수 있다면 적용하지 않는 것이 오히려 손해인 것처럼 보이기도 합니다. 

 

미국의 빅테크들이 이런 기술이 없어서 못만들었다기 보다는 CUDA를 포기해야하는 상황인데 CUDA가 갑자기 MLA를 지원합니다! 이렇게 되어버리면 빅테크들이 더이상 많은 그래픽카드가 필요하지 않아지고 이는 엔비디아의 실적 하락으로 이어지기 때문에 이걸 과연 적용해줄까 싶긴 합니다. 

 

그럼에도 미국 빅테크들도 결국 자신들의 비용을 줄이고 싶어할겁니다. 구글이 최근 TPU라는 칼을 빼들면서 엔비디아에 도전장을 내밀었었죠. 구글의 TPU가 추론에 특화된 AI칩인데 이게 구글의 전성비를 크게 높여주어 돈을 많이 아낄 수 있었습니다. 저는 개인적으로 구글이 새로운 커널을 만들어 MLA까지 적용하는 그림을 조심스럽게 바라고 있습니다. 

 

마치며

이번 포스팅에선 KV Cache를 기술 발전 순서에 맞춰서 다뤄보고 끝으로 제가 생각한 어른들의 사정까지 다뤄봤습니다. 이 내용을 공부할 때 수학 수식에 압도당해서 어떻게든 이해하려고 엄청나게 고생을 좀 했습니다. 이번 포스팅을 계기로 선형대수학을 공부해야하나.. 싶은 생각도 들었습니다. 

 

이번 포스팅을 위해 공부한 내용이 제 머리속에서 오랫동안 남아있기를 기원하면서.. (ㅜㅅㅜ) 이번 포스팅은 여기서 줄여보도록 하겠습니다. 긴 글 읽어주셔서 감사합니다. 오늘도 즐거운 하루 되세요!