Post

[LLM 추론 최적화] Day 3: KV 캐시와 Continuous Batching

[LLM 추론 최적화] Day 3: KV 캐시와 Continuous Batching

서론: Decode가 느린 이유

LLM Decode는 토큰을 하나씩 생성한다. 매 스텝마다 이전 모든 토큰의 Key-Value 쌍을 다시 읽어야 한다. 이것이 KV 캐시다. 문제는 이 캐시를 어떻게 관리하느냐에 따라 메모리 효율과 처리량이 크게 달라진다.

1. KV 캐시란

Transformer의 Attention 연산에서 매 토큰 생성 시 필요한 Key, Value 행렬을 저장한 것이다.

1
2
3
토큰 1 생성: K1, V1 계산 → 캐시 저장
토큰 2 생성: K1, V1 (캐시에서 읽기) + K2, V2 (새로 계산)
토큰 N 생성: K1..N-1, V1..N-1 (캐시) + KN, VN (새로 계산)

캐시 덕분에 이전 토큰을 재연산하지 않는다. 그 대신 시퀀스가 길어질수록 캐시 메모리가 선형으로 증가한다.

KV 캐시 크기 공식 (근사)

1
크기 = 2 × num_layers × num_heads × head_dim × seq_len × dtype_bytes

LLaMA-3 70B, FP16, 시퀀스 4096 토큰 기준 약 20GB.

2. 기존 정적 배치의 문제

1
2
3
4
5
6
7
8
정적 배치 방식:
  요청 1 (512토큰)  ─────────────────────■ (완료)
  요청 2 (2048토큰) ─────────────────────────────────■ (완료)
  요청 3 (대기 중)  ················· (요청 2 완료까지 대기)

문제:
  - 요청 1 완료 후 GPU 슬롯이 비어도 요청 3을 채울 수 없음
  - 짧은 요청과 긴 요청이 섞이면 GPU 낭비 심각

3. Continuous Batching

요청이 완료될 때마다 새 요청을 즉시 배치에 채워 넣는 방식이다. 배치를 고정하지 않고 iteration마다 재구성한다.

1
2
3
4
5
Continuous Batching:
  iter 1: [req1, req2, req3]
  iter 2: [req1, req2, req3]      ← req1 완료
  iter 3: [req4, req2, req3]      ← 즉시 req4 투입
  iter 4: [req4, req2, req5]      ← req3 완료, req5 투입

Orca 논문(2022)에서 처음 제안했고, vLLM이 이를 구현해 널리 알려졌다. GPU 활용률이 기존 정적 배치 대비 2-4배 개선된다고 알려져 있다.

4. PagedAttention

KV 캐시의 메모리 단편화 문제를 운영체제의 페이징 방식으로 해결한다.

4.1 기존 방식의 단편화

1
2
3
4
5
요청 A: 최대 2048 토큰 공간 예약 (실제 사용: 512토큰)
요청 B: 최대 2048 토큰 공간 예약 (실제 사용: 1800토큰)

→ 예약했지만 사용하지 않는 공간: 내부 단편화
→ 할당했다 해제한 빈 공간: 외부 단편화

4.2 PagedAttention의 접근

KV 캐시를 고정 크기 블록(페이지)으로 나누고, 논리적 시퀀스와 물리적 블록 간 매핑 테이블을 유지한다.

1
2
3
4
5
6
7
8
블록 크기: 16 토큰
요청 A (시퀀스 길이 37):
  논리 블록 0 → 물리 블록 7
  논리 블록 1 → 물리 블록 2
  논리 블록 2 → 물리 블록 9 (5토큰만 사용)

→ 내부 단편화: 마지막 블록의 일부만
→ 외부 단편화: 없음 (블록 단위 재사용 가능)

블록은 여러 요청이 공유할 수도 있다. 같은 시스템 프롬프트를 쓰는 요청들이 프리픽스 블록을 공유하면 KV 캐시 메모리를 크게 절약한다.

5. Prefix Caching

동일한 프롬프트 접두사를 가진 요청이 반복될 때, 해당 블록의 KV 캐시를 재사용한다.

1
2
3
4
적합한 시나리오:
  - 시스템 프롬프트가 긴 서비스 (RAG, 에이전트)
  - 멀티턴 대화 (이전 대화 KV 재사용)
  - 문서 분석 (같은 문서에 여러 질의)

긴 시스템 프롬프트(4K+ 토큰) 환경에서 TTFT를 50% 이상 줄일 수 있다.

6. Day 3 체크리스트

  1. 현재 추론 서버가 Continuous Batching을 지원하는지 확인했다.
  2. GPU 메모리 중 KV 캐시가 차지하는 비율을 측정했다.
  3. PagedAttention 기반 서버(vLLM 등)를 검토했다.
  4. 시스템 프롬프트가 긴 경우 Prefix Caching 적용 가능성을 확인했다.

다음 글 예고

Day 4에서는 추론 서버 설계를 다룬다. vLLM, TGI, NVIDIA Triton의 구조 차이, 선택 기준, 멀티 GPU 배포 전략을 정리한다.

This post is licensed under CC BY 4.0 by the author.