대규모 연산의 구원 투수, KV 캐시

앞의 글을 읽으면서 무언가 아찔함을 느끼셨다면 여러분은 감이 좋으신겁니다.

고작 단어 1개를 예측하려고 "기존 문장 전체를 다시 처음부터" 백여 개의 층을 거치며 매 층마다  N × N 어텐션 계산과 순전파 연산을 해야 한다? 그리고 그 과정을 또 최종 결과가 나올 때까지 반복한다? 뭔가 엄청 큰게 오는것 같은 느낌 아닌가요?

예를 들어 이 질문을 다시 들여다 봅시다. (간단하게 단어 하나가 토큰 하나라고 가정)

"사과에 가장 많이 들어 있는 비타민은?"

이 질문은 6개의 토큰이 있으니, 매번 어텐션 연산을 할 때마다 6x6개의 연산을 합니다. 이것을 GPT-3기준 96번 반복하죠. 그럼 6x6x96 이 됩니다.

그 결과 "사과에는" 이라는 단어 하나를 도출해 내죠. 그 다음에는 "사과에 가장 많이 들어 있는 비타민은?" 사과에는" 을 다시 입력으로 넣습니다. 이번에는 토큰이 7개로 늘었죠. 이번에는 7x7x96 개 연산을 합니다. 앞에 연산한것과 합하면 6x6x96 + 7x7x96 이 되죠. 이런식으로

"사과에 가장 많이 들어 있는 비타민은? ▶ 사과에는 비타민 C가 8~14% 정도로 가장 많이 들어 있습니다." 라는 답변을 완성할 때까지 6x6x96 + 7x7x96 .... + 15x15x96 번 연산을 해야 합니다.

6×6×96+7×7×96+⋯+15×15×96 =113,760

와우.. 이 간단한 질문과 답변을 하기 위해 11만번의 연산을 해야 하는군요.

질문과 답변이 이처럼 짧을 때는 그나마 괜찮습니다. 하지만 여러분이 수천 단어로 된 긴 논문을 주고 요약해 달라고 했다면 어떨까요? 단어 1개가 추가될 때마다 그 수천 개의 옛날 단어들을 전부 다시 처음부터 서로 연관성을 확인(어텐션)하는 엄청난 중노동을 해야 합니다. 이건 아무리 좋은 슈퍼컴퓨터라도 속도가 급격하게 느려질 수밖에 없는 방식입니다.

예를 들이 지금 쓰고 있는 이 글을 100단어로 요약하는 일을 시켰다고 생각해 봅시다.

이 글은 대략 400토큰 정도로 가정합시다. 이것에 대해 100토큰짜리 요약문을 생성한다면 앞의 수식처럼 400x400x96 + 401x401x96 + 402x402x96... 499x499x96 이런식으로 계산되서 19억번 정도의 어텐션 연산이 필요하게 됩니다.

(이것은 이해하기 쉽게 설명한 것이고, 실제로는 트랜스포머 한번 당 NxN 연산은 두번 일어 납니다. 어텐션 점수 만들때 한번 컨텍스트 벡터 만들 때 한번. 이외에도 다른 연산이 많기 때문에 실제로는 이 숫자의 몇배의 연산이 필요합니다.)

어마어마하죠? 만약 이 상태였다면 GPT는 세상에 출시되지 못했을 겁니다.

그래서 이런 비효율 문제를 해결하기 위해 엔지니어들은 KV 캐시(Key-Value Cache)라는 방법을 만들었습니다.

모델은 단어 하나를 생성할 때마다 수십 개의 트랜스포머 계층을 거치며 각 단어의 문맥 정보(Key와 Value)를 계산합니다. 그리고 이렇게 계산한 결과를 버리지 않고 임시 메모리 공간(Cache)에 저장해 둡니다.

덕분에 다음 단어를 생성할 때는 이전 문장을 처음부터 다시 계산할 필요가 없습니다. 새롭게 추가된 단어 하나에 대해서만 연산을 수행하고, 과거 단어들과의 관계를 계산할 때는 캐시에 저장해 두었던 이전 계산 결과를 그대로 가져다 사용하면 됩니다.

즉, 이미 계산한 문맥 정보는 재사용하고, 새로 들어온 단어만 추가로 계산하는 방식입니다. 그래서 문장이 길어져도 훨씬 빠르게 답변을 생성할 수 있습니다.

위 그림을 보면 좀더 이해하기 쉬우실 겁니다.

덕분에 모델은 매번 기존 문장 전체를 처음부터 다시 계산하는 끔찍한 연산 낭비를 피할 수 있습니다.

400토큰짜리 이 글을 요약하는 앞에 예로든 사례에 KV캐시가 적용되면 19억번 계산 하는 대신 다음과 같이 1,967만번만 계산하는 것으로 확 줄어듭니다.

400x400x96 + 400x96 + 401x96 + 499x96 = 약 1,967만번

(처음 400x400x96만 NxN연산이고, 그 다음부터는 새로 들어온 단어에 대해서만 연산)

대화가 길어져도 GPT가 이전 내용을 매번 처음부터 다시 계산하지 않고 빠르게 응답할 수 있는 이유가 바로 이 KV 캐시 덕분입니다.

최근 HBM을 생산하는 한국 기업들 주가가 천정부지로 치솟는 이유 중 하나도 이 KV캐시 수요가 늘어서 그렇습니다. 

1명의 사용자가 8000토큰 분량의 대화를 내누면 약 1.2GB의 KV캐시가 필요하다고 합니다. 100만명이 동시에 사용한다면 1.2PB (페타바이트)의 메모리가 필요하구요. 요즘은 컨텍스트 윈도우의 크기가 점점 더 커지고 있고 사용자들도 더 늘어나고 있어서 메모리 수요는 한동안 계속 높을것 같습니다. (컨텍스트 윈도우는 다음장에 설명)