KV Cache · STUDY

Autoregressive generation은 새 token을 붙일 때마다 길어진 prefix를 조건으로 다음 token을 고른다. 그러나 이전 prefix 전체의 states를 매번 다시 계산할 필요는 없다. Causal Mask 때문에 과거 position은 새로 붙은 미래 token을 보지 않았고, 따라서 그 과거 K/V도 변하지 않기 때문이다. KV Cache는 이 재사용 가능성을 이용하되, 계산 대신 요청별 memory를 더 사용한다.

Prompt Prefill
cached $K$ / $V$
new token $Q$ / $K$ / $V$
append $K$ / $V$
new $Q$ attends prefix
Decode 반복

Reusing Past K/V

Causal decoding의 과거 state는 미래 token을 보지 않았으므로 과거 $K$ / $V$ 가 변하지 않아 재사용할 수 있다. Cache 없이는 step t의 prefix Attention이 $O(t^{2})$ , 있으면 새 query의 $O(t)$ 다. $n$ -token 생성 전체 Attention work는 $O(n^{3})$ → $O(n^{2})$ 다.

새 step에서는 새 position의 Q/K/V만 계산하고, K/V를 기존 cache 뒤에 붙인다. 새로운 Query는 과거 Keys와 compatibility를 구하고 그에 대응하는 Values를 읽는다. 이전 Queries의 outputs는 새 output을 계산하기 위해 다시 만들 필요가 없으므로, 저장하여 다시 읽는 대상은 K/V다.

K_{1:t}= \begin{bmatrix}K_{1:t-1}\\k_t\end{bmatrix}, \qquad V_{1:t}= \begin{bmatrix}V_{1:t-1}\\v_t\end{bmatrix}

이 동작은 layer마다 자신의 cache에 대해 수행한다. Cache가 없으면 모든 prefix Queries의 scores를 다시 구하지만, 있으면 이번 Query row만 계산한다. 아래 $O(n^3)\to O(n^2)$ 는 이 one-layer Attention score work를 생성 전체에 합한 비교이며, model 전체 latency의 정확한 비율은 아니다.

Cache는 model weights가 아니라 request-local inference state다. 다른 요청은 다른 prefix를 가지므로 weights처럼 모든 요청에 동일한 값 하나를 공유한다고 생각해서는 안 된다.

Prefill & Decode

2,000-token prompt / 300-token answer에서 Prefill은 prompt 전체 처리, Decode는 새 token 반복이다. TTFT는 queue+Prefill+첫 생성, TPOT는 연속 output tokens 사이의 latency다. $T_{\mathrm{response}}\approx\mathrm{TTFT}+(N_{\mathrm{out}}-1)\mathrm{TPOT}$ .

Prompt 길이는 Prefill뿐 아니라 길어진 $K$ / $V$ 를 읽는 Decode에도 영향을 줄 수 있다.

Memory Footprint

Llama 2 7B style MHA, 32 layers, 32 KV heads, $d_{head}=128$ , BF16, no tensor parallel 가정이다.

1 Token / 1 Layer:

2\times128\times32\times2\,\mathrm{B}=16\,\mathrm{KiB}

All Layers:

32\text{ layers}=512\,\mathrm{KiB/token}

One Request:

4096\text{ tokens}=2\,\mathrm{GiB}

Concurrent Requests:

32\text{ concurrent requests}=64\,\mathrm{GiB}

BF16 model weights≈13GiB보다 cache가 커질 수 있다. 8 KV heads GQA는 동일 조건 cache의 $\frac14$ 이다.

Latency and Capacity

Prefill은 2,000-token prompt가 이미 주어진 상태에서 전체를 처리해 cache를 만든다. 그 뒤 300-token answer의 Decode는 앞에서 고른 token에 의존하여 한 번씩 진행한다. TTFT에는 대기 시간, Prefill, 첫 token 생성이 들어가고, TPOT는 이후 두 output tokens 사이의 시간이다. 그래서 300 tokens라면 첫 token 이후 간격은 299개다.

Prompt가 길면 TTFT뿐 아니라 TPOT도 영향을 받을 수 있다. Decode 때 더 긴 K/V cache를 계속 읽기 때문이다. Batch를 늘리면 weights를 읽는 비용을 여러 요청에 활용할 수 있지만, 동시에 유지할 cache도 늘어난다. 두 latency 지표와 serving memory를 한 가지 원인으로만 나누지 않는다.

예시의 16 KiB/token/layer는 K와 V, 32 heads, 128 coordinates, BF16 2 bytes를 곱한 값이다. 이를 32 layers, 4,096 tokens, 32 requests로 차례로 늘리면 64 GiB가 된다. Model weights 약 13 GiB보다 요청별 cache가 더 클 수 있다는 것이 중요하다. KV heads를 32에서 8로 줄이는 GQA는 같은 조건에서 cache를 1/4로 줄인다. 다음 section에서 Query 표현과 KV 공유가 어떻게 함께 유지되는지 본다.

M_{\mathrm{KV}}=2BLnH_{\mathrm{KV}}d_{\mathrm{head}}\times\mathrm{bytes}

T_{\mathrm{response}}\approx\mathrm{TTFT}+(N_{\mathrm{out}}-1)\mathrm{TPOT}

\text{Attention work: generation of }n\text{ tokens}

\text{No cache:}\quad\sum_tO(t^2)=O(n^3)

\text{Cache:}\quad\sum_tO(t)=O(n^2)

KV Cache는 모델 weights가 아니라 요청별 state다. FLOPs 절약과 저장 memory 증가를 함께 본다.

요청마다 cache가 있고 live tokens·동시 요청 수에 비례한다.

Training · Prefill · Decode

Known sequence의 병렬 처리와 아직 모르는 다음 token 생성을 구분한다.

Aspect	Training	Prefill	Decode
입력	Known sequence / gold targets	Prompt 전체	새로 생성한 token
계산	Positions parallel + mask	Prompt positions 처리	이전 생성에 의존하는 sequential steps
저장/신호	Valid target CE / backward	Prompt $K$ / $V$ cache 생성	새 $K$ / $V$ 추가 / prefix cache 읽기
latency 관점	학습 처리량	TTFT에 큰 영향	TPOT에 큰 영향

TTFT에는 queueing/first generation도 포함되고 긴 cache는 Decode latency에도 영향을 준다.

Check

과거 $Q$ 를 매 step 재사용하지 않고 $K$ / $V$ 를 cache하는 이유는?

Answer

현재 새 token의 Query가 과거 Keys/Values를 읽는다. 과거 Query output을 다시 쓸 필요는 없고 causal 구조에서 과거 $K$ / $V$ 는 변하지 않는다.

32 layers,32 KV heads,width128,BF16의 4096-token request32개에 필요한 cache를 계산하라.

Answer

2( $K$ , $V$ )×32×32×128×2B=512KiB/token. 4096 tokens는 2GiB/request,32개면 64GiB다.

Cache를 썼는데 긴 prompt에서 Decode도 느려질 수 있는 이유는?

Answer

새 $Q$ 가 길어진 prefix $K$ / $V$ 를 읽어야 한다. Weights/KV memory access, batch, 동시 요청 조건도 TPOT에 영향을 준다.

새 token이 추가되어도 이전 K/V를 다시 만들지 않아도 되는 조건은 무엇인가?

Answer

Causal model에서는 과거 position이 새로 추가된 미래 token을 읽지 않으므로 과거 states와 K/V가 바뀌지 않는다. 새 position의 K/V를 append하고 Query가 전체 prefix cache를 읽는다. 재계산을 줄이는 대신 요청별 cache memory가 늘어난다.

Prefill 후 새 token에서 KV Cache가 어떻게 갱신되고 읽히는가?

Answer

새 token의 $Q$ / $K$ / $V$ 계산 → 새 $K$ / $V$ 를 layer별 cache에 append → 새 $Q$ 가 cached $K$ / $V$ 전체 prefix를 읽음 → output token 선택 → 다음 step 반복.

Reusing Past K/V

Prefill & Decode

Memory Footprint

Latency and Capacity

Training · Prefill · Decode

Check

Related