KV 캐시 완전정복: LLM 추론 속도와 메모리 비용을 결정하는 트랜스포머 핵심 구조 이해하기

핵심 메시지

• 자기회귀 생성은 매 단계마다 앞선 모든 토큰을 다시 처리하면 답 전체로는 n의 세제곱 규모의 낭비가 생기는데, KV 캐시가 이 반복 계산을 없앤다.
• 각 토큰의 Q(질의)는 그 토큰 출력을 계산하는 순간 한 번만 쓰이고 버려지지만, K와 V는 미래의 모든 토큰이 계속 참조하므로 K·V만 캐시한다.
• 캐시 덕분에 단계당 연산은 시퀀스 길이에 선형으로 줄지만, 캐시 자체가 사용자당 수십 GB에 이르는 메모리·대역폭 비용으로 되돌아온다.
• 디코딩은 연산보다 메모리가 병목(산술 강도 약 1 flop/byte)이라 GPU 코어가 데이터를 기다리며 놀고, 그래서 양자화가 추론을 극적으로 빠르게 한다.
• 프롬프트(프리픽스) 캐시는 처음부터 연속된 접두부에만 유효하므로, 고정된 내용을 앞에 두고 가변 내용을 뒤에 두는 배치가 성능을 좌우한다.

쉽게 이해하기

언어모델이 답을 절반쯤 썼다고 하자. 다음 단어를 만들려면 앞선 모든 단어를 다시 보고 전체 어텐션 층을 돌려야 한다. 문제는 그다음 단어를 만들 때 한 토큰 더 길어진 시퀀스에 대해 이 과정을 처음부터 또 반복한다는 점이다. 답이 1,000 토큰이면 프롬프트의 첫 단어는 약 1,000번 다시 처리된다. 단계당 n 제곱, 답 전체로는 n 세제곱의 연산인데, 정작 앞부분 단어들은 전혀 바뀌지 않는다.

해결의 실마리는 어텐션 내부의 비대칭에 있다. 모든 토큰은 Q(무엇을 찾는가), K(무엇을 제공하는가), V(실제로 기여하는 값)로 나뉜다. 그런데 한 토큰의 Q는 그 토큰의 출력이 계산되는 그 순간에만 쓰이고 이후엔 필요 없다. 반면 K와 V는 미래의 모든 토큰이 매 디코딩 단계마다 조회한다. 그래서 규칙은 저절로 도출된다. 영원히 재사용되는 K·V는 캐시하고, 다시 쓰이지 않는 Q는 버린다.

이 캐시는 공짜가 아니다. 하나의 텐서가 아니라 모든 어텐션 층, 모든 KV 헤드, 현재 문맥의 모든 토큰에 대한 K와 V다. 총 바이트는 2(K·V) × 층 수 × KV 헤드 수 × 시퀀스 길이 × 헤드 차원 × 원소당 바이트다. 예컨대 Llama 3 70B를 반정밀도로 쓰면 4,000 토큰 문맥에 약 2.5GB, 32,000 토큰이면 사용자당 무려 20GB의 캐시가 모델 가중치 위에 얹힌다. 긴 문맥 서빙이 어렵다는 말의 실체가 바로 이 캐시다.

여기에 마스킹의 세부가 아름답게 맞물린다. 학습·프리필 단계에서는 시퀀스를 병렬 처리하되 앞 토큰이 미래 토큰을 보지 못하게 막아야 한다. 마스크는 Q·K·V가 아니라 소프트맥스 직전의 점수 행렬에 얹힌다. 미래를 보는 위치에 음의 무한대를 더하면 소프트맥스 뒤 정확히 0이 된다. 0을 곱하지 않고 음의 무한대를 더하는 이유는, 소프트맥스에서 점수 0은 e^0=1이라 오히려 정상적인 한 표가 되기 때문이다.

디코딩이 빠른 하드웨어에서도 느린 이유는 메모리 병목이다. 70B 모델 한 단계는 약 140GB의 가중치를 고대역폭 메모리에서 불러오면서 약 1,400억 번의 유용한 연산을 한다. 바이트와 연산이 비슷해 산술 강도가 약 1 flop/byte인데, H100의 균형점은 약 300 flop/byte라 코어가 데이터를 기다리며 논다. 그래서 양자화(연산 수는 같고 바이트만 줄임)가 큰 도움이 되고, 배칭은 모두가 공유하는 가중치에는 도움이 되지만 사용자별로 따로인 캐시 대역폭에는 그렇지 않다. 마지막으로 프리픽스 캐시는 층이 깊어질수록 각 토큰의 K·V가 앞 토큰 전체와 얽히므로, 접두부의 한 바이트만 달라도 그 지점부터 캐시가 무효가 된다. 고정 내용은 앞, 가변 내용은 뒤라는 배치 원칙이 여기서 나온다.

주요 인사이트

• '비활성화'를 곱셈 0이 아니라 덧셈 음의 무한대로 구현하는 이유가 핵심이다. 소프트맥스에서 점수 0은 e^0=1이 되어 여전히 유효한 표가 되지만, 음의 무한대는 e^(-∞)=0으로 그 위치를 합에서 완전히 제거한다.
• 디코딩 성능의 본질은 연산이 아니라 메모리 대역폭이다. 산술 강도 약 1 대 H100 균형점 약 300이라는 격차가 GPU가 노는 이유를 설명하며, 이 때문에 양자화가 추론 속도에 크게 기여한다.
• 배칭은 사용자 전체가 공유하는 가중치 로딩만 분산시킬 뿐, 사용자마다 별도인 KV 캐시의 대역폭 부담은 배치 크기에 그대로 비례해 늘어난다.
• 층이 깊어질수록 한 토큰의 K·V는 앞선 모든 토큰과 깊게 얽힌다. 그래서 프롬프트 재사용은 처음부터 연속된 접두부에만 가능하고, 중간이나 접미부는 재사용할 수 없다.
• 다중 질의 어텐션, 슬라이딩 윈도, 프롬프트 캐싱 같은 현대 추론 최적화는 결국 캐시를 줄이거나(shrink), 공유하거나(share), 캐시 로드당 연산을 늘리는(extract) 세 갈래로 수렴한다.

자주 묻는 질문

관련 AI 소식