Paged Attention: Prefill vs Decode

1. LLM Inference의 두 단계: Prefill vs Decode

LLM이 텍스트를 생성할 때는 반드시 두 가지 단계를 거칩니다. 이 두 단계의 차이를 정확히 이해하는 것이 KV Cache를 이해하는 핵심입니다.

1-1. Prefill 단계 (Prompt Processing)

정의: 사용자가 입력한 전체 프롬프트(문장)를 한 번에 모델에 넣어, 각 토큰의 Key와 Value를 한 번에 계산하는 단계입니다.

왜 필요한가: LLM은 문맥을 이해해야 합니다. 입력 문장 전체를 한 번에 처리하여 각 토큰 간의 관계를 파악하고, 이 정보를 KV Cache에 저장해두어야 합니다.

1-2. Decode 단계 (Token Generation)

정의: Prefill이 끝난 후, 한 토큰씩 순차적으로 새로운 토큰을 생성하는 단계입니다.

왜 필요한가: LLM은 자동회귀(Autoregressive) 모델입니다. 즉, 이전에 생성한 토큰을 바탕으로 다음 토큰을 예측합니다. 이때 이전 토큰들의 정보(K, V)를 매번 다시 계산하지 않고 Cache에서 불러와야 효율적입니다.

---

2. 예시로 상세히 풀어보기

예시 시나리오

사용자 입력: "나는 학교에"
모델이 생성해야 할 출력: "간다" → "." → <EOS> (종료 토큰)

토큰 목록: ["나는", "학교에", "간다", ".", "<EOS>"]

임베딩 차원 $d_{model} = 4$, Head 차원 $d_k = d_v = 2$ (이전 단계와 동일)

헤드 차원 ($d_k = d_v = 2$)

트랜스포머의 핵심은 '멀티 헤드 어텐션(Multi-Head Attention)'이야. 이건 "여러 명의 전문가(Head)가 단어를 나누어서 분석한다"는 뜻입니다.

여기서 $d_k = 2$라는 건, 전체 특징(4개)을 전문가들이 2개씩 나누어서 본다는 뜻입니다.

전문가들의 역할 분담 예시

전체 특징이 [특징1, 특징2, 특징3, 특징4]라면:

1. 전문가 A (Head 1): 앞의 2개 숫자 [특징1, 특징2]만 집중해서 봐. (예: "이 단어가 사람인지 장소인지 위주로 판단해!")
2. 전문가 B (Head 2): 뒤의 2개 숫자 [특징3, 특징4]만 집중해서 봐. (예: "이 단어가 동작인지 문법적인 건지 위주로 판단해!")

• 계산 공식: $d_{model} (4) \div \text{헤드 수} (2) = d_k (2)$
• 즉, 이 모델은 2명의 전문가(Head)가 있고, 각 전문가가 2차원($d_k=2$)의 정보를 담당하는 구조입니다.

---

3. Prefill 단계 상세 풀이

Step 1: 입력 토큰 임베딩

Prefill 단계에서는 사용자가 입력한 "나는"과 "학교에"를 처리합니다.

위치	토큰	임베딩 벡터 (4차원)
1	나는	$[1, 0, 1, 0]$
2	학교에	$[0, 1, 1, 0]$

Step 2: Self-Attention 계산 (이전 단계와 동일)

이전 단계에서 계산한 대로:

$Q_{prefill} = \begin{bmatrix} 2 & 0 \\ 1 & 1 \end{bmatrix}, \quad K_{prefill} = \begin{bmatrix} 0 & 2 \\ 1 & 1 \end{bmatrix}, \quad V_{prefill} = \begin{bmatrix} 1 & 2 \\ 1 & 1 \end{bmatrix}$

Attention Output:
$\text{Output}_{prefill} \approx \begin{bmatrix} 1.768 & 1.813 \\ 1.504 & 1.752 \end{bmatrix}$

Step 3: KV Cache에 저장

Prefill이 끝나면, 이 $K_{prefill}$과 $V_{prefill}$을 메모리(KV Cache)에 저장합니다.

$\text{KV Cache} = \left\{ K = \begin{bmatrix} 0 & 2 \\ 1 & 1 \end{bmatrix}, \quad V = \begin{bmatrix} 1 & 2 \\ 1 & 1 \end{bmatrix} \right\}$

이 캐시는 이후 Decode 단계에서 계속 사용됩니다.

---

4. Decode 단계 상세 풀이 (KV Cache의 핵심)

이제 모델은 "간다"라는 새로운 토큰을 생성해야 합니다.

4-1. KV Cache가 없다면? (비효율적인 경우)

"간다"를 생성하기 위해 "나는", "학교에", "간다" 전체에 대해 Self-Attention을 다시 계산해야 합니다.

새로운 토큰 "간다"의 임베딩: $[1, 1, 0, 1]$

새로운 $Q, K, V$ 행렬 (3×2):
$Q_{new} = \begin{bmatrix} 2 & 0 \\ 1 & 1 \\ 1 & 2 \end{bmatrix}, \quad K_{new} = \begin{bmatrix} 0 & 2 \\ 1 & 1 \\ 2 & 1 \end{bmatrix}, \quad V_{new} = \begin{bmatrix} 1 & 2 \\ 1 & 1 \\ 2 & 2 \end{bmatrix}$

Attention Score (3×3):
$Q_{new}K_{new}^T = \begin{bmatrix} 0 & 2 & 4 \\ 2 & 2 & 3 \\ 4 & 3 & 4 \end{bmatrix}$

→ "나는"과 "학교에"의 K, V를 이미 Prefill에서 계산했는데, 다시 계산하는 중복 발생!

4-2. KV Cache가 있다면? (효율적인 경우)

"간다" 토큰만 새로 계산하고, 이전 토큰들의 K, V는 Cache에서 불러옵니다.

새로 계산해야 할 것 (오직 "간다" 토큰만):

• $Q_{간다} = [1, 1, 0, 1] \times W_Q = [1, 2]$
• $K_{간다} = [1, 1, 0, 1] \times W_K = [2, 1]$
• $V_{간다} = [1, 1, 0, 1] \times W_V = [2, 2]$

Cache에서 불러올 것:

• $K_{나는} = [0, 2]$, $V_{나는} = [1, 2]$
• $K_{학교에} = [1, 1]$, $V_{학교에} = [1, 1]$

확장된 KV Cache:
$K_{cache} = \begin{bmatrix} 0 & 2 \\ 1 & 1 \\ 2 & 1 \end{bmatrix}, \quad V_{cache} = \begin{bmatrix} 1 & 2 \\ 1 & 1 \\ 2 & 2 \end{bmatrix}$

"간다" 토큰의 Attention Score 계산:

$Q_{간다} = [1, 2]$와 모든 $K$의 내적:

• $Q_{간다} \cdot K_{나는} = 1\times0 + 2\times2 = 4$
• $Q_{간다} \cdot K_{학교에} = 1\times1 + 2\times1 = 3$
• $Q_{간다} \cdot K_{간다} = 1\times2 + 2\times1 = 4$

$\text{Score} = [4, 3, 4]$

Softmax 적용 ($\sqrt{2} \approx 1.414$로 나눈 후):

• $[4/1.414, 3/1.414, 4/1.414] = [2.828, 2.121, 2.828]$
• $e^{2.828} \approx 16.92$, $e^{2.121} \approx 8.34$
• 합계: $16.92 + 8.34 + 16.92 = 42.18$
• Softmax: $[16.92/42.18, 8.34/42.18, 16.92/42.18] = [0.401, 0.198, 0.401]$

"간다" 토큰의 Attention Output:
$\text{Output}_{간다} = 0.401 \times [1, 2] + 0.198 \times [1, 1] + 0.401 \times [2, 2]$
$= [0.401 + 0.198 + 0.802, \; 0.802 + 0.198 + 0.802]$
$= [1.401, \; 1.802]$

---

5. 연산량 비교: KV Cache 유무

구분	KV Cache 없음	KV Cache 있음
"간다" 생성 시 필요한 K 계산	3개 토큰 모두 재계산	1개(새 토큰)만 계산
"간다" 생성 시 필요한 V 계산	3개 토큰 모두 재계산	1개(새 토큰)만 계산
Attention 행렬 크기	3×3	1×3 (Query 1개 × Key 3개)
FLOPs (부동소수점 연산)	$O(n^2 \cdot d)$	$O(n \cdot d)$ (Query-Key 내적만)

여기서 $n$은 시퀀스 길이, $d$는 차원입니다.

시퀀스 길이가 1000개, 2000개로 길어질수록 KV Cache의 효율성은 기하급수적으로 증가합니다.

---

6. KV Cache의 메모리 구조 시각화

---

7. 시퀀스 길이에 따른 연산량 차이 (구체적 수치)

예시로 $d_k = 64$, 배치 크기 = 1, Layer = 1이라고 가정할 때:

시퀀스 길이 $n$	KV Cache 없음 (FLOPs)	KV Cache 있음 (FLOPs)	절약률
10	$\approx 12,800$	$\approx 1,280$	90%
100	$\approx 1,280,000$	$\approx 12,800$	99%
1,000	$\approx 128,000,000$	$\approx 128,000$	99.9%
4,096 (LLaMA)	$\approx 2.1 \times 10^9$	$\approx 524,288$	99.97%

→ 시퀀스가 길어질수록 KV Cache는 거의 필수가 됩니다.

---

8. KV Cache의 메모리 구조 (행렬 관점)

실제로 Transformer는 여러 Layer와 여러 Attention Head를 가집니다.

단일 Layer, 단일 Head의 KV Cache

$\text{K Cache} \in \mathbb{R}^{n \times d_k}, \quad \text{V Cache} \in \mathbb{R}^{n \times d_v}$

실제 LLM의 KV Cache (예: LLaMA-2 7B)

• Layer 수: 32
• Head 수 per Layer: 32
• $d_k$ per Head: 128
• $d_{model}$: 4,096

한 토큰당 KV Cache 메모리:
$\text{Memory per token} = 2 \times (\text{num\_layers}) \times (\text{num\_heads}) \times (d_k) \times (\text{bytes})$
$= 2 \times 32 \times 32 \times 128 \times 2 \text{ bytes (FP16)}$
$= 524,288 \text{ bytes} \approx 0.5 \text{ MB per token}$

4,096 토큰의 문맥을 가진다면:
$4,096 \times 0.5 \text{ MB} = 2,048 \text{ MB} \approx 2 \text{ GB}$

→ 단일 요청당 2GB의 KV Cache가 필요합니다.

---

📌 2단계 정리

개념	정의	왜 사용하는가
Prefill	입력 프롬프트 전체를 한 번에 Self-Attention 처리	문맥 파악 및 초기 KV Cache 생성
Decode	한 토큰씩 순차적으로 생성	자동회귀적 텍스트 생성
KV Cache	이전 토큰들의 K, V 행렬을 메모리에 저장	중복 연산 제거, O(n²) → O(n)으로 개선
Cache Update	새 토큰의 K, V를 기존 Cache에 append	다음 Decode 단계에서 재사용 가능

---

⚠️ KV Cache의 한계 (Paged Attention으로 가는 다리)

지금까지 KV Cache가 얼마나 효율적인지 살펴봤습니다. 하지만 실제 서빙 환경에서는 다음과 같은 심각한 문제가 발생합니다:

1. 메모리 낭비 (Internal Fragmentation): 각 요청의 시퀀스 길이가 다르므로, 미리 할당한 고정 크기의 KV Cache 공간이 낭비됨
2. 외부 단편화 (External Fragmentation): 요청들이 메모리에 불연속적으로 배치되어 큰 연속 메모리를 할당하지 못함
3. 공유 불가: 동일한 프롬프트를 여러 요청이 공유할 수 없음 (Beam Search, Parallel Sampling 등)

이 문제들을 해결하기 위해 등장한 것이 바로 Paged Attention입니다.