KV 캐시(KV Cache) 완전 이해: LLM 추론을 빠르게 하는 어텐션 최적화 원리

핵심 메시지

• LLM은 토큰을 하나씩 생성하며, 매 단계 이전 토큰 전체를 다시 계산해 계산량이 이차로 증가한다.
• KV 캐시는 각 단계에서 계산한 키(K)·값(V) 표현을 저장해 두고 다음 단계에서 재사용한다.
• 캐시를 쓰면 새로 생성된 마지막 토큰의 쿼리만 계산하면 되어 어텐션 계산이 이차에서 선형으로 줄어든다.
• 작은 모델 기준 Apple M3 CPU에서 약 7~8배의 생성 속도 향상이 확인됐다.
• 대신 시퀀스 길이에 비례해 메모리 사용량이 크게 늘어나는 트레이드오프가 있다.

쉽게 이해하기

LLM은 텍스트를 토큰 단위로 하나씩 만들어 내며, 매 단계에서 지금까지 생성된 모든 토큰에 대해 계산을 반복한다. 시퀀스가 길어질수록 계산량이 제곱으로 불어나는데도 현대 LLM이 빠르게 동작하는 핵심 이유가 바로 KV 캐시다.

어텐션에서는 각 층마다 모든 토큰에 대해 쿼리(Q)·키(K)·값(V)을 만든다. 한 토큰의 쿼리와 모든 이전 토큰의 키를 내적해 어텐션 가중치를 얻고, 그 가중치로 값들을 가중합해 문맥 표현을 만든다. 문제는 생성 과정에서 이전 토큰들의 K·V가 변하지 않는데도 단계마다 똑같이 다시 계산된다는 점이다.

해결책은 단순하다. 각 단계에서 계산한 K·V를 캐시에 저장해 두고, 다음 단계에서는 마지막 토큰의 Q·K·V만 새로 계산한 뒤 이전 토큰들의 K·V는 캐시에서 꺼내 쓴다. 다음 토큰 예측에는 마지막 토큰의 쿼리만 필요하므로 이전 토큰들의 쿼리는 저장할 필요조차 없다. 첫 단계에서는 프롬프트 전체를 처리해 캐시를 채우고, 이후에는 새로 나온 토큰 하나만 처리한다.

효과는 분명하다. 캐시가 없으면 모든 쿼리와 키 사이의 N×N 행렬을 계산하지만, 캐시를 쓰면 마지막 행 하나만 계산하면 되어 어텐션 계산량이 시퀀스 길이에 대해 이차에서 선형으로 줄어든다. 영상은 minGPT 구현에 K·V 캐시 텐서를 미리 0으로 할당하고, 유효 길이를 추적하는 변수와 use_cache 플래그, 1×N 마스크, 절대 위치 임베딩의 위치 보정, 생성 시작 전 캐시 초기화 등을 더해 이를 코드로 구현한다.

결과적으로 캐시가 없으면 토큰이 길어질수록 토큰당 시간이 계속 늘지만, 캐시를 쓰면 거의 일정하게 유지되어 작은 모델·Apple M3 CPU 기준 약 7~8배 빨라졌다. 다만 대가가 있다. 캐시 메모리는 층 수·헤드 수·헤드 차원·시퀀스 길이에 비례해 선형으로 증가하며, 80층·8헤드·헤드차원 128·문맥 8K·FP16에 배치 크기 4를 적용하면 약 10GB에 이른다(모델 가중치와 활성값은 별도). 후속 영상에서는 이 메모리를 줄이는 기법을 다룬다고 예고한다.

주요 인사이트

• 생성 단계마다 변하지 않는 키·값을 다시 계산하던 낭비를 제거하는 것이 KV 캐시의 본질이다.
• 다음 토큰 예측에는 마지막 토큰의 쿼리만 필요하므로 이전 토큰들의 쿼리는 캐시에 저장하지 않는다.
• 계산량을 이차에서 선형으로 낮추는 대신, 캐시 메모리는 시퀀스 길이에 비례해 늘어나는 트레이드오프가 있다.
• 절대 위치 임베딩을 쓰면 캐시 모드에서 위치 인덱스를 누적 처리 길이로 따로 추적해야 정확하다.

자주 묻는 질문

관련 AI 소식