KV Cache

Transformer 추론에서 KV Cache가 필요한 이유와 메모리 사용량이 커지는 이유는 무엇인가.

Transformer 추론에서 KV Cache는 왜 필요하고, 컨텍스트가 길어질수록 메모리를 많이 소비하는 이유는 무엇인가?

autoregressive decoding에서 과거 토큰의 attention key/value를 매 토큰마다 다시 계산하면 연산 낭비가 크다.
KV Cache는 과거 토큰의 K/V tensor를 저장해 새 토큰의 attention 계산에서 재사용한다.
사용량은 레이어 수, KV head 수, head dimension, 저장 precision, sequence length, batch size 등에 비례한다.
Multi-Query Attention과 Grouped-Query Attention은 K/V head 수를 줄여 cache 사용량을 줄일 수 있다.

Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention: https://arxiv.org/abs/2309.06180
Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints: https://arxiv.org/abs/2305.13245

KV Cache는 Transformer 기반 LLM이 텍스트를 한 토큰씩 생성할 때 이전 토큰에서 계산한 attention의 Key와 Value 값을 저장해 재사용하는 메모리 영역이다.

이 캐시가 없으면 새 토큰 하나를 생성할 때마다 이전 모든 토큰의 Key와 Value를 다시 계산해야 하므로 추론이 매우 비효율적이다.

Autoregressive LLM은 다음 토큰을 생성할 때 지금까지 생성된 모든 토큰을 문맥으로 사용한다.

첫 번째 토큰에서 계산한 attention 정보는 두 번째 토큰에서도 필요하고, 세 번째 토큰에서도 다시 필요하다.

KV Cache는 이전 토큰의 K와 V를 저장해두고 새 토큰에 해당하는 값만 추가한다. 덕분에 decoding 단계의 반복 계산을 크게 줄일 수 있다.

KV Cache는 모델 가중치와 달리 요청마다 별도로 생긴다.

사용량은 대략 다음 요소에 비례한다.

즉 모델 가중치 크기가 같더라도 컨텍스트 길이를 늘리거나 동시 사용자 수를 늘리면 KV Cache 메모리가 빠르게 증가한다.

컨텍스트가 2배 길어지면 보통 저장해야 하는 과거 토큰의 K/V도 2배 가까이 늘어난다.

그래서 장문 모델의 실제 서빙에서는 '모델이 최대 몇 토큰을 지원하는가'와 '그 길이로 몇 요청을 동시에 처리할 수 있는가'가 서로 다른 문제다.

기본 Multi-Head Attention에서는 query마다 대응하는 여러 key/value head가 존재한다.

Multi-Query Attention은 여러 query head가 하나의 K/V를 공유하고, Grouped-Query Attention은 여러 query head가 그룹별 K/V를 공유한다.

K/V head 수가 줄어들면 KV Cache 크기와 memory bandwidth 요구량을 크게 줄일 수 있어 현대 LLM에서 널리 사용된다.

여러 요청을 동시에 처리하면 각 sequence의 길이가 달라 메모리 단편화가 생기기 쉽다.

PagedAttention은 KV Cache를 고정된 block 단위로 관리해 필요한 부분만 할당하고 재사용할 수 있게 하는 접근이다. vLLM이 이 아이디어를 LLM serving에 널리 알렸다.

가능한 방법에는 다음이 있다.

각 방법은 속도, 품질, 구현복잡도 사이의 trade-off가 있다.

KV Cache는 LLM decoding을 빠르게 만드는 대신 요청별 메모리를 소비한다. 긴 컨텍스트와 높은 동시성을 동시에 원하는 serving 시스템에서는 모델 가중치보다 KV Cache가 실제 수용량을 결정하는 병목이 될 수 있다.