[LLM 시간 및 메모리 최적화] 2. KV Cache & Paged Attention

KV Cache

Transformer 디코더의 각 레이어에서 $i$번째 Query, Key, Value의 행은 직전 레이어의 $i$번째 Query와 $1$부터 $i$번째 Key, Value의 행에 영향을 받는다.

다시 말해, $i$번째 Query, Key, Value의 행은 직전 레이어의 $i$번째 이후의 Query, Key, Value의 행에 전혀 영향을 받지 않는다.
$i$시점에 계산된 Query, Key, Value는 그 이후로도 변하지 않는다.

그렇기 때문에, $i+1$번째 토큰을 예측할 때, 이전에 계산했던 $1$부터 $i$번째 Key, Value의 행을 그대로 재활용할 수 있다는 것이다.

(참고로, $i+1$번째 토큰 예측 시 오직 $i+1$번째 Query 행만 필요하기 때문에, Query에서 재활용할 수 있는 정보가 없다.)

 

KV Cache 기법은 이전에 계산한 $1$부터 $i-1$번째 Key, Value의 행을 HBM(or MEM)에 저장하여, $i$번째 토큰 예측 시 이를 사용해 Attention 연산을 수행함으로써, 불필요한 Key, Value의 반복 연산 방지하는 기법이다.

출처: https://medium.com/@joaolages/kv-caching-explained-276520203249

Paged Attention

기존 KV Cache 기법은 메모리를 최대 토큰 길이(max seq len)만큼 할당(reservation)한 다음 연속적(contiguous)으로 KV를 저장하고 있다. 이 방식은 실제 활용되는 토큰 길이에 비해 과도한 메모리를 할당하기 때문에, 전체 할당된 메모리의 20~40%만 사용(utilize)되는 등 내부 단편화(internal memory fragmentation)가 심각하다.

 

Paged Attention은 KV Cache를 연속적(contiguous)으로 저장하는 대신, 블록(block) 단위로 비연속적(non-contiguous)으로 관리하여 할당된 메모리의 활용도를 94%까지 높이며, 내부 단편화 문제를 효과적으로 해결한 기법이다.

Memory sharing

동일한 입력 프롬프트에 대해 여러 요청이 들어와 여러 개의 output을 생성할 때, 입력 프롬프트의 KV cache를 공유해 메모리를 절약할 수 있다.