개요
Autoregression모델에서, KV Cache는 전에 계산한 K값과 V값들을 캐싱하여, 현재 토큰을 계산하는데 필요한 중복되는 연산을 줄여서 생성 속도의 향상을 꾀하는 Caching모델이다. KV Cache는 대표적인 Computing & Memory trade-off를 이용한 방법이다.
필요성
Autoregressive모델에서 디코더가 매 타임스텝마다 다시 연산을 수행하는 것은 비효율적이다. KV Cache는 이전 타임스텝에서 계산된 K, V 값을 저장하여 이후 연산에서 다시 활용하는 방식으로 동작한다. KV Cache를 사용하면 이전 단어들의 Key/Value값을 재사용하여 연산량을 줄이고, 추론 속도를 개선 할수 있다.
그림에서 알수 있다 싶이[1], Regression모델의 계산은 QKV값을 계산해야 한다. 각 디코딩 스텝마다 새로운 Key와 Value값을 계산하기 때문에 결과적으로 알고리즘의 복잡도는 [math]\displaystyle{ O(n^2) }[/math]이 된다. 그러나 KV Cache를 사용하면 새로운 Query만 계산하면 됨으로 연산량이 [math]\displaystyle{ O(n) }[/math]로 감소하는 효과를 얻을 수 있다.
한계
KV Cache는 이전 토큰들의 Key Value값을 메모리에 계속 저장해야 함으로 메모리 사용량이 증가한다. 특히 긴 문장을 처리하는 경우에, GPU VRAM사용량이 급격하게 늘어날 수 있다. 예를 들어서 KV Cache의 메모리 요구량은 다음과 같다.
Transformer 모델에서 KV Cache는 각 Attention Head마다 Key (K)와 Value (V) 행렬을 저장해야 한다. 따라서, 전체 KV Cache의 메모리 요구량은 다음과 같이 계산된다.
- [math]\displaystyle{ \text{Memory}_{KV} = 2 \times B \times H \times N \times S \times d_k \times \text{size}(\text{dtype}) }[/math]
- B : 배치 크기 (Batch Size)
- H : 헤드 개수 (Number of Attention Heads)
- S : 현재까지 입력된 토큰 개수 (Sequence Length)
- N : Decoder에서 사용된 레이어의 개수
- [math]\displaystyle{ d_k }[/math] : 각 Attention Head의 Key/Value 차원 (Key/Value Dimension)
- [math]\displaystyle{ \text{size}(\text{dtype}) }[/math] : 데이터 타입(Float16, Float32 등)에 따른 단일 요소 크기 (예: Float16이면 2바이트, Float32이면 4바이트)
- 앞의 2배는 Key와 Value 두 개의 행렬을 저장하기 때문이다.
즉, KV Cache의 메모리 사용량은 배치 크기(B), 시퀀스 길이(L), 헤드 개수(H), Key/Value 차원(d_k)에 선형적으로 증가한다.
- 시퀀스 길이 S가 증가하면 메모리 사용량도 비례하여 증가 → 긴 문맥을 다룰수록 더 많은 메모리 필요.
- 헤드 개수 H 가 많을수록 메모리 사용량 증가 → 모델 크기가 클수록 더 많은 캐시 필요.
- [math]\displaystyle{ d_k }[/math] 가 증가할수록 메모리 사용량 증가 → 더 높은 차원의 Key/Value 벡터를 사용하면 성능이 향상되지만 메모리도 더 필요.
예를 들어서 Llama3-70B모델을 생각해보면, Layer의 크기가 80, 헤드 개수 64개 헤드당 차원이 128개, 최대 생성가능한 Sequence길이가 4k인 모델이기 때문에, KV Cache의 크기는 다음과 같다.
- [math]\displaystyle{ \text{Memory}_{KV} = 2 \times 1 \times 64 \times 80 \times 128 \times 4096 \times 2 }[/math]
로 계산되어서, 대략 토큰당 ~2.5MB 최대 시퀀스 길이 생성시 10.5GB의 KV Cache크기가 소요된다. 만약 여러 클라이언트가 동시에 요청하여서 Batch size가 커지면, KV Cache의 크기또한 기하급수적으로 커진다.