KV-cache는 LLM 추론(inference)에서 가장 중요한 최적화 기법이며, 모델과의 실시간 채팅이 가능한 이유입니다. 이것이 무엇이고 왜 중요한지 알아보겠습니다.

생성 과정은 자기회귀적(autoregressive)입니다

LLM은 토큰을 하나씩 생성합니다. 토큰을 방출하고, 이를 추가하며, 다음 토큰을 위해 전체 모델을 다시 실행하는 방식입니다. 각 어텐션 레이어 내부에서 모든 토큰은 Query, Key, Value가 됩니다. 가장 새로운 토큰을 생성하기 위해서는 그 Query가 이전의 모든 토큰들의 Key와 점수화되고, 이 가중치들이 그 Value들을 혼합합니다. 따라서 토큰 $t$를 생성하려면 토큰 1부터 $t$까지의 K와 V가 필요합니다.

순진한 접근 방식은 이차적(quadratic)입니다

캐시가 없다면, 매 단계마다 전체 프롬프트 접두사(prefix)를 재인코딩하여 토큰 1부터 $t$까지의 K/V를 다시 구축해야 합니다. 1단계는 1개의 토큰을 처리하고, 2단계는 2개를, ..., N단계는 N개를 처리합니다. 총 작업량은 $\approx 1+2+\dots+N = \mathbf{N(N+1)/2}$이며, 이는 이차적($\text{quadratic}$)입니다. 토큰 1의 K/V는 변하지 않음에도 불구하고 매 단계마다 재계산됩니다.

핵심 통찰: 과거의 K/V는 절대 변하지 않습니다

LLM은 인과 마스킹(causal masking)을 사용합니다. 즉, 한 토큰이 오직 이전 토큰들만 참조한다는 의미입니다. 따라서 끝에 새로운 토큰을 추가한다고 해서 이전 토큰들의 Key와 Value가 바뀔 수는 없습니다. 이들은 상수입니다. 이를 재계산하는 것은 순수한 낭비입니다.

캐싱 → 선형적(linear) 생성

각 토큰의 K/V를 처음 계산할 때 저장합니다. 각 단계에서는 오직 하나 개의 새로운 토큰에 대한 K/V만 계산하고, 이를 추가한 다음 전체 캐시를 참조합니다:

K_cache, V_cache = [], []
for t in range(N):
    k, v = kv(token_t)              # 하나의 토큰 작업량
...

단계별 작업량이 일정해지므로 총 작업량은 $O(N^2)$ 대신 $\mathbf{O(N)}$이 됩니다.

Prefill 대 Decode

이는 추론 과정을 두 단계로 나눕니다. Prefill: 전체 프롬프트를 한 번의 병렬 패스로 받아 캐시를 채웁니다. 계산량이 많기 때문에 긴 프롬프트에서 첫 토큰을 생성하는 데 시간이 걸릴 수 있습니다. Decode: 출력 토큰을 하나씩 생성하며, 이는 저렴한 캐시 추가 작업입니다. 이것이

캐시는 모든 토큰, 모든 레이어(layer), 모든 헤드(head)에 대한 K와 V를 저장합니다. 그 크기는 컨텍스트 길이(context length)에 따라 선형적으로 증가하며, 이것이 바로 128k-토큰 컨텍스트가 비용이 많이 드는 정확한 이유입니다. 캐시는 수 기가바이트의 GPU 메모리를 점유할 수 있으며, 종종 하나의 GPU가 얼마나 많은 사용자를 서비스할 수 있는지를 결정하는 제한 요소가 됩니다. paged attention (vLLM), grouped-query attention, 양자화된 캐시(quantised caches), 그리고 프롬프트 캐싱(prompt caching)과 같은 기법들은 모두 이를 제어하기 위해 존재합니다.

"캐시가 없는" 생성과 "캐시가 있는" 생성이 토큰별로 어떻게 달라지는지 확인해 보세요:
https://dev48v.infy.uk/ai/days/day22-kv-cache.html