KV Cache

inferenceservingvllmattentionmemory

KV cache는 LLM 추론에서 가장 기본적인 최적화다.

한 문장으로 말하면:

이미 처리한 token의 Key와 Value를 layer마다 저장해 둔다.

그래서 다음 token을 만들 때 과거 token 전체를 다시 계산하지 않아도 된다.

왜 필요한가

autoregressive generation은 token을 한 개씩 만든다.

prompt -> token 1
prompt + token 1 -> token 2
prompt + token 1 + token 2 -> token 3
...

KV cache가 없다면 매 step마다 prefix 전체를 다시 transformer에 넣어야 한다.

예를 들어 이미 1,000 token을 읽었고 다음 token을 하나 만들고 싶다고 하자. cache가 없으면 모델은 1,001 token 전체에 대해 layer 계산을 다시 한다.

이건 대부분 중복 계산이다. 과거 1,000 token의 K와 V는 이미 이전 step에서 계산한 값과 같다.

무엇을 저장하는가

attention에서 각 token은 세 가지 벡터를 만든다.

Q: 지금 무엇을 찾고 싶은가
K: 이 token은 어떤 주소/특징을 갖는가
V: 이 token에서 실제로 가져올 정보는 무엇인가

decode step에서 다음 token을 예측하는 데 필요한 query는 새 token의 Q 하나다. 과거 token들의 Q는 다시 필요하지 않다.

하지만 새 Q는 과거 모든 token의 K와 비교해야 하고, attention weight로 과거 모든 V를 섞어야 한다.

그래서 cache에는 Q가 아니라 K와 V를 저장한다.

저장한다: past K, past V
저장하지 않는다: past Q

Decode 한 step에서 일어나는 일

KV cache가 있으면 decode step은 이렇게 바뀐다.

1. 새 token 하나에 대해서만 Q, K, V를 계산한다.
2. 새 K, V를 KV cache 뒤에 append한다.
3. 새 Q가 cache에 있는 모든 K를 본다.
4. attention 결과로 cache에 있는 모든 V를 섞는다.
5. 다음 token logits를 만든다.

이제 transformer는 매 step마다 prefix 전체가 아니라 새 token 하나만 처리한다. 중복 계산이 크게 줄어든다.

무엇이 사라지고 무엇이 남는가

KV cache가 없애는 것은 과거 token의 K/V 재계산이다.

하지만 attention 자체의 과거 참조가 사라지는 것은 아니다. 새 token은 여전히 과거 모든 token을 봐야 한다.

없어지는 것:
과거 token의 Q/K/V, MLP 계산을 매번 다시 하는 일

남는 것:
새 Q가 모든 과거 K/V를 읽고 attention하는 일

그래서 KV cache는 decode를 빠르게 만들지만, 동시에 새로운 병목을 만든다. 매 output token마다 GPU는 model weight와 KV cache를 계속 HBM에서 읽어야 한다.

왜 memory-bound가 되는가

prefill에서는 많은 token을 한 번에 처리하기 때문에 큰 행렬 곱셈이 생긴다. weight를 읽어도 여러 token 계산에 재사용할 수 있다.

decode에서는 새 token이 보통 하나뿐이다.

prefill: many tokens -> GEMM에 가까움 -> weight 재사용이 큼
decode : one token   -> GEMV에 가까움 -> weight 재사용이 작음

KV cache 덕분에 계산량은 줄었지만, 매 step마다 읽어야 하는 데이터는 여전히 크다.

매 decode step에서 읽는 것:
- model weights
- past tokens의 KV cache

짧은 context에서는 model weight 읽기가 더 크게 보일 수 있다. 긴 context나 큰 batch에서는 KV cache 읽기와 저장 공간이 점점 중요해진다.

vLLM에서 왜 중요한가

vLLM을 이해하려면 KV cache를 단순한 “속도 최적화”가 아니라 serving system의 핵심 자원으로 봐야 한다.

요청이 많아질수록 각 요청은 자기만의 KV cache를 가진다. output이 길어질수록 cache도 길어진다. GPU memory는 제한되어 있다.

따라서 serving engine은 다음 질문을 계속 풀어야 한다.

어떤 요청의 KV cache를 GPU memory에 둘 것인가?
새 token의 K/V를 어디에 붙일 것인가?
요청이 끝나면 cache 공간을 어떻게 회수할 것인가?
서로 길이가 다른 요청들을 어떻게 batch로 묶을 것인가?

이 질문들이 PagedAttention과 scheduler로 이어진다.

KV cache는 vLLM의 메모리 관리자, block table, PagedAttention kernel을 읽기 위한 출발점이다.

연결

확인

  • KV cache에는 왜 Q가 아니라 K와 V를 저장하는가?
  • KV cache가 없으면 decode step마다 어떤 중복 계산이 생기는가?
  • KV cache가 attention의 전체 과거 참조 비용을 완전히 없애지 못하는 이유는 무엇인가?
  • vLLM에서 KV cache가 단순 tensor가 아니라 관리해야 할 serving 자원인 이유는 무엇인가?