KV 캐시

KV 캐시(Key-Value Cache)는 트랜스포머 기반 오토리그레시브 언어 모델에서, 이전 토큰의 정보를 저장하여 추론 시 반복 계산을 줄이기 위해 사용하는 캐시 구조이다.

개념

• 트랜스포머 디코더는 각 층에서 self-attention을 계산할 때, 이전 토큰의 key와 value 벡터를 반복적으로 참조한다.
• 이를 효율화하기 위해 한 번 계산한 key와 value를 캐시에 저장해 두고, 이후 토큰 생성 시 다시 사용한다.
• 이때 저장된 key/value 쌍들을 "KV 캐시"라고 부른다.

동작 방식

• 디코더는 입력 시퀀스를 한 토큰씩 생성하며, 각 스텝에서 이전 토큰의 K/V를 활용해 현재 토큰의 attention을 계산한다.
• 추론 시 각 스텝에서 새로 생성된 토큰에 대한 key와 value만 새로 추가되고, 기존 KV 캐시는 그대로 재사용된다.
• 전체 문맥 길이에 대해 linear하게 캐시가 커진다.

구성

• 각 디코더 층마다 KV 캐시를 따로 유지한다.
• KV 캐시는 일반적으로 다음 구조로 구성된다:
  • K: [batch_size, num_heads, seq_len, head_dim]
  • V: [batch_size, num_heads, seq_len, head_dim]
• seq_len은 지금까지 생성된 토큰 수(=캐시 길이)이며, 생성이 진행될수록 증가한다.

장점

• 연산 효율성: 이미 계산된 key/value를 재사용하므로 중복 계산을 줄인다.
• 추론 속도 향상: 토큰 생성이 점진적으로 이루어지는 오토리그레시브 구조에서 속도 개선 효과가 크다.
• 긴 문맥 처리 가능: 길어진 문맥을 한 번에 계산하는 대신, 캐시로 분할 저장하며 처리할 수 있다.

한계

• 메모리 사용 증가: 문장이 길어질수록 KV 캐시도 커지며, 특히 대형 모델에서는 메모리 사용량이 급증한다.
• 지연 시간(latency): GPU/TPU 메모리에서의 캐시 관리, 복사 비용 등이 지연을 유발할 수 있다.
• 구조적 유연성 부족: 고정된 attention 구조에서는 캐시가 비효율적으로 확장될 수 있다.

최적화 기법

• KV 캐시 가지치기: 중요하지 않은 K/V 벡터를 제거하여 메모리 절감
• 캐시 압축(Cache Compression): float16, int8 등 저정밀도로 저장
• Cross-layer 공유: 일부 계층의 K/V만 유지하거나 공유하여 메모리 최적화
• Lazy KV 캐시: 반드시 필요한 시점에만 K/V를 생성하거나 저장

같이 보기

• 트랜스포머 (신경망)
• 오토리그레시브 모델
• KV 캐시 가지치기
• 시퀀스-투-시퀀스 모델