KV 캐시 완전정리: 트랜스포머 추론이 GPU 메모리를 많이 쓰는 이유와 속도 최적화 원리

핵심 메시지

• 트랜스포머는 텍스트를 길게 생성할수록 GPU 메모리를 더 많이 써서 결국 메모리 부족으로 멈출 수 있다.
• 추론 중 메모리의 상당 부분은 키·값을 저장하는 KV 캐시가 차지한다.
• 자기회귀 생성에서 이전 토큰의 키·값 행렬은 거의 바뀌지 않으므로 매번 다시 계산하는 것은 낭비다.
• KV 캐시는 이전 키·값을 저장해 새 토큰마다 열·행 하나씩만 계산하도록 만들어 연산을 크게 줄인다.
• KV 캐시 덕분에 토큰당 작업량이 시퀀스 길이와 무관하게 일정해진다.

쉽게 이해하기

트랜스포머 기반 언어 모델의 근본적 한계 하나는 텍스트를 더 많이 생성할수록 GPU 메모리를 더 많이 소모한다는 점이다. 결국 메모리가 바닥나면 프로그램이 멈춰 더 이상 생성할 수 없다. 긴 문맥 모델의 입력 토큰 요금이 더 비싼 것도 이런 높은 메모리 사용의 경제적 결과로 볼 수 있다.

셀프 어텐션을 짚어 보면, 각 토큰의 임베딩 벡터 X에 학습된 세 행렬을 곱해 쿼리·키·값을 만든다. 디코딩 중에는 새 토큰 하나를 나타내는 쿼리가 벡터인 반면, 이전 문맥 전체를 담는 키와 값은 행렬이다. 어텐션은 쿼리와 키의 내적에 소프트맥스를 취해 값 행렬에 가중합으로 적용한다.

한 번에 한 단어씩 생성하는 자기회귀 방식에서 키·값 행렬은 시퀀스가 길어져도 크게 바뀌지 않는다. 한 단어의 임베딩이 한 번 계산되면 이후로는 변하지 않기 때문이다. 그런데도 매 단계마다 모든 단어의 키·값을 다시 계산하면 행렬-벡터 곱이 제곱으로 늘어 매우 느려진다. 마치 다음 단어를 쓸 때마다 앞서 쓴 모든 단어를 다시 읽는 것과 같다.

KV 캐시는 이 낭비를 없앤다. 새 단어가 들어오면 쿼리 벡터는 새로 만들되, 이전 키·값은 저장해 두었다가 재사용하고 새 토큰에 해당하는 키 행렬의 열 하나와 값 행렬의 행 하나만 계산한다. 셀프 어텐션 층은 이전 KV 캐시와 현재 토큰 임베딩을 받아 새 키·값을 만들어 캐시에 덧붙여 GPU 메모리에 저장한다. 현재 토큰과 이전 문맥이 상호작용하는 곳은 셀프 어텐션 층뿐이므로, 새 토큰마다 필요한 작업량은 일정하게 유지된다.

메모리 사용량은 키·값 두 행렬, 파라미터당 바이트 수(정밀도), 층 수, 임베딩 차원, 시퀀스 길이, 배치 크기를 모두 곱해 구한다. 300억 파라미터급 모델을 16비트로 추론하고 층 48개, 임베딩 차원 약 7,000, 배치 128을 가정한 예에서 KV 캐시는 약 180GB로, 모델 자체(약 60GB)의 세 배에 이른다. 또한 프롬프트를 처음 처리할 때는 캐시가 없어 모든 토큰의 키·값을 계산해야 하므로 지연이 크고, 이후 토큰은 하나씩만 계산하므로 지연이 낮아진다.

주요 인사이트

• 추론 시 메모리 병목은 모델 가중치가 아니라 KV 캐시인 경우가 흔하다.
• 키·값이 한 번 정해지면 바뀌지 않는다는 성질이 캐싱을 가능하게 하는 핵심이다.
• KV 캐시로 토큰당 연산이 상수가 되어, 긴 시퀀스에서도 작업량이 급증하지 않는다.
• KV 캐시 크기는 층 수·차원·시퀀스 길이·배치 크기에 비례하므로 긴 문맥일수록 급격히 커진다.
• 프롬프트 처리(첫 토큰)의 지연이 이후 토큰보다 큰 이유는 캐시가 아직 비어 있기 때문이다.

자주 묻는 질문

관련 AI 소식