[Daily morning study] LLM 추론 최적화: KV 캐시와 투기적 디코딩
#daily morning study
LLM 추론의 기본 구조
Transformer 기반 LLM은 토큰을 하나씩 순차적으로 생성하는 자기회귀(Autoregressive) 방식으로 동작한다.
입력: "The capital of France is"
출력: "Paris" → " The" → " Eiffel" → ...
각 토큰을 생성할 때마다 이전 모든 토큰에 대한 Attention 연산을 수행한다. 입력이 길수록, 생성이 길어질수록 연산량이 기하급수적으로 증가한다.
추론 과정은 크게 두 단계로 나뉜다:
| 단계 | 이름 | 설명 |
|---|---|---|
| 1단계 | Prefill | 입력 프롬프트 전체를 한 번에 처리 (병렬 연산) |
| 2단계 | Decode | 토큰을 하나씩 생성 (순차 연산) |
Prefill은 GPU를 잘 활용하지만, Decode는 토큰 하나를 생성할 때마다 전체 모델을 통과하기 때문에 느리다. 추론 최적화의 핵심 목표는 Decode 단계를 빠르게 만드는 것이다.
KV 캐시(KV Cache)
왜 필요한가
Transformer의 Self-Attention에서 각 레이어는 세 가지 행렬을 계산한다: Query(Q), Key(K), Value(V).
Attention(Q, K, V) = softmax(QK^T / √d_k) · V
토큰 i를 생성할 때, 이전 토큰들(0~i-1)의 K, V 값은 이미 계산된 값과 동일하다. 그런데 매번 토큰을 생성할 때마다 이 값을 다시 계산하면 엄청난 낭비다.
KV 캐시는 이미 계산한 K, V 행렬을 메모리에 저장해두고, 새 토큰 생성 시 재사용하는 기법이다.
동작 방식
# 캐시 없이 (비효율)
1번째 토큰 생성: K[0], V[0] 계산
2번째 토큰 생성: K[0], K[1], V[0], V[1] 모두 재계산
3번째 토큰 생성: K[0..2], V[0..2] 모두 재계산
...
# KV 캐시 사용 (효율)
1번째 토큰 생성: K[0], V[0] 계산 → 캐시에 저장
2번째 토큰 생성: K[1], V[1]만 계산 → 캐시에 append
3번째 토큰 생성: K[2], V[2]만 계산 → 캐시에 append
새 토큰에 해당하는 Q만 계산하고, 저장된 K, V 캐시와 Attention을 수행한다.
KV 캐시 메모리 계산
KV 캐시가 차지하는 메모리는 생각보다 크다.
KV 캐시 크기 = 2 × num_layers × num_heads × head_dim × seq_len × batch_size × dtype_size
# LLaMA-2 7B 기준 (FP16):
# layers=32, heads=32, head_dim=128, dtype_size=2bytes
# seq_len=4096, batch_size=1
= 2 × 32 × 32 × 128 × 4096 × 1 × 2 bytes
≈ 2 GB
배치 크기나 시퀀스 길이가 늘어나면 KV 캐시 메모리도 비례해서 증가한다. 메모리 부족이 곧 추론 처리량(throughput)의 병목이 된다.
Paged Attention (PagedAttention)
기존 KV 캐시의 문제점
기존 KV 캐시는 연속된(contiguous) 메모리 블록에 시퀀스별로 할당된다. 문제는:
- 각 요청의 최대 시퀀스 길이를 미리 알 수 없음
- 최대 길이 기준으로 선점 할당하면 실제 사용량보다 훨씬 많은 메모리 낭비
- 메모리 단편화(fragmentation) 발생
vLLM이 제안한 PagedAttention은 OS의 가상 메모리(paging) 아이디어를 LLM 추론에 적용했다.
동작 방식
물리 메모리를 고정 크기 블록(block)으로 나눔
각 요청의 KV 캐시는 비연속적인 블록들에 저장 가능
블록 테이블(block table)이 논리 블록 → 물리 블록 매핑 관리
| 비교 항목 | 기존 방식 | PagedAttention |
|---|---|---|
| 메모리 할당 | 연속 블록, 최대 길이 선점 | 비연속 블록, 필요할 때 동적 할당 |
| 메모리 낭비 | 20~40% 낭비 | 거의 없음 (<4%) |
| GPU 활용도 | 낮음 | 높음 |
| 처리량 | 낮음 | 2~4배 높음 |
PagedAttention 덕분에 같은 GPU 메모리로 더 많은 요청을 동시에 처리할 수 있다.
투기적 디코딩(Speculative Decoding)
아이디어
KV 캐시는 중복 연산을 줄이지만, 결국 토큰을 하나씩 생성하는 순차적 특성은 변하지 않는다. 투기적 디코딩은 이 병목을 완전히 다른 방식으로 공략한다.
핵심 아이디어: 작은 드래프트 모델(draft model)로 여러 토큰을 빠르게 추측하고, 큰 타깃 모델(target model)이 한 번에 검증하면 빠르지 않을까?
동작 과정
1. 드래프트 모델(small): 토큰 k개를 빠르게 추측
→ "Paris", "is", "the", "capital" (4토큰 추측)
2. 타깃 모델(large): 추측된 4토큰 전체를 병렬로 검증
→ 타깃 모델도 사실 4개 토큰에 대해 한 번의 포워드 패스로 처리 가능
3. 검증 결과 처리:
- 모두 일치: 4개 토큰 수용 (4배 속도 향상)
- n번째 토큰 불일치: n-1개 수용 + 타깃 모델의 토큰으로 교체
드래프트 모델이 틀려도 타깃 모델의 출력 품질은 보장된다. 틀린 경우에는 그냥 타깃 모델을 한 번 돌린 것과 같은 비용만 소모된다.
수락률(Acceptance Rate)과 실제 효과
평균 수락률 α라 할 때:
기대 토큰 수 per step ≈ (1 - α^k) / (1 - α) (k: 드래프트 길이)
α = 0.8, k = 5이면:
≈ (1 - 0.8^5) / (1 - 0.8) = (1 - 0.328) / 0.2 ≈ 3.36 토큰/스텝
실제로는 2~3배 속도 향상이 보고되며, 특히 반복적인 패턴이 많은 코드 생성이나 번역에서 효과가 크다.
드래프트 모델 선택 전략
| 방식 | 설명 |
|---|---|
| 별도 소형 모델 | LLaMA-7B를 LLaMA-70B의 드래프트로 사용 |
| 같은 모델의 레이어 스킵 | 타깃 모델의 앞쪽 레이어만 사용 (Medusa, EAGLE) |
| n-gram 예측 | 단순 텍스트 패턴 매칭으로 드래프트 생성 |
연속 배치(Continuous Batching)
전통적인 배치 처리는 모든 요청이 끝날 때까지 기다린 후 다음 배치를 처리한다. 짧은 요청이 먼저 끝나도 긴 요청이 완료될 때까지 GPU는 낭비된다.
연속 배치는 하나의 요청이 완료되면 즉시 새 요청을 채워넣는다.
# 전통적 배치
[요청A ████████████] [요청B ██] [요청C ██████] ← B, C가 끝나도 A 대기
새 요청 → 다음 배치 대기
# 연속 배치 (iteration-level scheduling)
[요청A ████████████]
[요청B ██][새 요청D ██████] ← B 완료 즉시 D 투입
[요청C ██████]
vLLM, TGI(Text Generation Inference) 같은 서빙 프레임워크들이 PagedAttention + Continuous Batching을 함께 구현하여 GPU 활용률을 크게 높인다.
Flash Attention
Attention 연산은 시퀀스 길이 N에 대해 O(N²)의 메모리를 사용한다. 긴 컨텍스트에서는 GPU의 HBM(High Bandwidth Memory)에 거대한 Attention 행렬을 올려야 한다.
Flash Attention은 Attention 행렬 전체를 메모리에 올리지 않고, SRAM(빠른 온칩 메모리)에서 타일 단위로 처리하는 IO-aware 알고리즘이다.
메모리 복잡도: O(N²) → O(N)
속도: 2~4배 빠름 (메모리 접근 감소)
출력: 수학적으로 정확히 동일 (근사 아님)
Flash Attention 2, 3으로 발전하면서 현재는 사실상 모든 대형 모델 학습과 추론의 표준이 되었다.
정리
| 기법 | 목표 | 핵심 아이디어 |
|---|---|---|
| KV Cache | 중복 연산 제거 | 이전 K,V 행렬 재사용 |
| PagedAttention | 메모리 낭비 제거 | OS paging 적용 |
| Speculative Decoding | Decode 병렬화 | 드래프트 추측 후 일괄 검증 |
| Continuous Batching | GPU 활용률 향상 | 완료 즉시 새 요청 투입 |
| Flash Attention | 메모리 접근 최소화 | 타일 기반 온칩 연산 |
이 기법들은 독립적이 아니라 서로 결합해서 사용된다. vLLM은 PagedAttention + Continuous Batching + Flash Attention을 통합해 단일 GPU에서 수십 배 높은 처리량을 달성한다. LLM 추론의 병목을 이해하면 왜 이 기법들이 등장했는지 자연스럽게 납득된다.
