Post

[Transformer 원리] Day 5: 추론 최적화 - KV 캐시, 양자화, FlashAttention

[Transformer 원리] Day 5: 추론 최적화 - KV 캐시, 양자화, FlashAttention

이 글은 AI(Claude)의 도움을 받아 작성하고, 작성자가 검토·편집했습니다.

서론: 원리에서 실전 서빙으로

Day 1~4에서 Transformer가 어떻게 텍스트를 처리하고 학습·생성하는지 완성했다. 하지만 수십억~수천억 파라미터 모델을 실제 서비스로 띄우는 순간, 속도와 메모리의 벽에 부딪힌다. 마지막 편은 이 벽을 넘는 추론 최적화 기법들 — KV 캐시, 양자화, FlashAttention, 그리고 효율적 서빙이다. 대규모 시스템 시리즈(Day 1~5)의 운영 관점과도 이어진다.

1. 추론의 두 단계와 병목

1
2
3
4
5
6
7
8
1. Prefill(프리필): 입력 프롬프트 전체를 한 번에 처리
   → 병렬화 잘 됨, 연산 집약적(compute-bound)

2. Decode(디코드): 토큰을 하나씩 생성 (자기회귀)
   → 한 번에 한 토큰, 메모리 대역폭 병목(memory-bound)
   → 토큰마다 전체 모델 가중치를 메모리에서 읽어야 함

핵심: LLM 추론 비용의 대부분은 "한 토큰씩 생성"하는 디코드 단계

2. KV 캐시: 가장 중요한 최적화

자기회귀 생성에서 매 토큰마다 이전 토큰들의 Key·Value를 다시 계산하면 막대한 낭비다. 이미 계산한 K·V를 저장(캐시)한다.

1
2
3
4
캐시 없이: 100번째 토큰 생성 시 1~99번 토큰의 K,V를 매번 재계산
캐시 있이: 1~99번의 K,V는 저장해 두고, 새 토큰의 K,V만 추가 계산

→ 생성 속도가 수십 배 빨라짐 (필수 최적화)
1
2
3
4
5
6
7
8
# KV 캐시의 개념
def decode_step(new_token, kv_cache):
    q, k, v = compute_qkv(new_token)        # 새 토큰만 계산
    kv_cache.k = concat(kv_cache.k, k)      # 캐시에 누적
    kv_cache.v = concat(kv_cache.v, v)
    # 새 Query로 누적된 전체 K,V에 어텐션
    output = attention(q, kv_cache.k, kv_cache.v)
    return output, kv_cache

대가는 메모리다. KV 캐시는 (레이어 × 헤드 × 시퀀스 길이 × 차원)에 비례해, 긴 컨텍스트·다중 사용자에서 메모리를 크게 잡아먹는다. 이를 줄이는 것이 다음 기법들이다.

3. KV 캐시 줄이기: MQA / GQA

1
2
3
4
MHA (Multi-Head): 헤드마다 독립 K,V → KV 캐시 큼 (Day 3)
MQA (Multi-Query): 모든 헤드가 K,V를 공유 → 캐시 대폭 감소, 품질 약간 저하
GQA (Grouped-Query): 헤드를 그룹으로 묶어 그룹당 K,V 공유
  → MHA 품질과 MQA 효율의 절충 (현대 LLM 다수가 채택)

4. 양자화: 정밀도를 낮춰 메모리 줄이기

가중치를 낮은 비트로 표현해 메모리와 대역폭을 줄인다.

1
2
3
4
5
6
7
FP16/BF16 (16비트): 학습·기본 추론의 표준
INT8 (8비트): 메모리 절반, 속도↑, 품질 손실 미미
INT4 (4비트): 메모리 1/4, 소비자 GPU에서도 대형 모델 구동 가능

기법:
  PTQ(학습 후 양자화): 학습된 모델을 사후 변환 (GPTQ, AWQ)
  QAT(양자화 인지 학습): 학습 중 양자화 시뮬레이션 → 품질 더 좋음
1
2
3
4
예: 70B 모델 메모리
  FP16: ~140GB (고가 GPU 여러 장 필요)
  INT4: ~35GB  (단일 고급 GPU로 구동 가능)
→ 양자화가 LLM의 접근성을 결정적으로 넓혔다

5. FlashAttention: 메모리를 똑똑하게 쓰기

표준 어텐션은 [n×n] 점수 행렬(Day 2)을 통째로 메모리에 만든다. 긴 시퀀스에서 이게 폭발한다.

1
2
3
4
5
6
7
8
9
표준 어텐션 문제:
  n×n 행렬을 GPU의 느린 메모리(HBM)에 읽고 씀
  → 시퀀스 길이의 제곱으로 메모리·대역폭 증가

FlashAttention:
  행렬을 작은 블록(tile)으로 나눠, 빠른 메모리(SRAM) 안에서 처리
  전체 n×n 행렬을 메모리에 저장하지 않고 점진적으로 계산
  ✅ 메모리 사용 선형화, 2~4배 속도, 긴 컨텍스트 가능
  (결과는 수학적으로 동일 — 근사가 아니라 정확한 재배치)

핵심 통찰: 연산량이 아니라 메모리 접근이 병목이므로, 메모리 이동을 줄이는 것이 답이다(eBPF Day 4의 “어디서 시간을 쓰는가”와 같은 사고).

6. 처리량 높이기: 배칭과 서빙

단일 요청 속도뿐 아니라 전체 처리량(throughput)도 중요하다(대규모 시스템 Day 3).

1
2
3
4
5
6
7
8
9
10
11
연속 배칭(Continuous Batching):
  여러 요청을 함께 처리하되, 끝난 요청은 빼고 새 요청을 즉시 채움
  → GPU를 놀리지 않음, 처리량 대폭 향상 (vLLM의 핵심)

PagedAttention:
  KV 캐시를 OS의 페이징처럼 블록 단위로 관리
  → 메모리 단편화 제거, 더 많은 동시 요청 수용

투기적 디코딩(Speculative Decoding):
  작은 모델이 여러 토큰을 미리 추측 → 큰 모델이 한 번에 검증
  → 품질 유지하며 속도 향상

이런 기법들은 vLLM, TensorRT-LLM, TGI 같은 추론 서버에 구현돼 있어, 직접 만들기보다 활용하는 것이 정석이다.

7. 시리즈 종합 체크리스트

  1. 토큰화·임베딩·위치 인코딩으로 텍스트를 벡터로 바꿨다. (Day 1)
  2. Q·K·V 셀프 어텐션으로 토큰이 문맥을 읽는 원리를 이해했다. (Day 2)
  3. 멀티헤드·FFN·잔차·정규화로 Transformer 블록을 조립했다. (Day 3)
  4. 사전학습·정렬·디코딩 전략으로 모델이 배우고 생성하는 법을 파악했다. (Day 4)
  5. KV 캐시·양자화·FlashAttention·배칭으로 추론을 최적화했다. (Day 5)

시리즈 마무리

Transformer의 위력은 단순한 구성 요소의 반복에서 나온다. 어텐션으로 정보를 모으고, FFN으로 변환하고, 잔차로 깊이 쌓고, “다음 토큰 예측”이라는 단일 과제로 학습한다. 이 단순함이 규모 확장을 가능케 했고, 그 결과가 오늘의 LLM이다.

토큰화→어텐션→블록→학습/추론→최적화 다섯 단계를 따라오면, LLM이 더 이상 블랙박스가 아니다. “왜 컨텍스트 길이가 비싼가(KV 캐시)”, “왜 temperature가 창의성을 바꾸나(샘플링)”, “왜 양자화가 품질을 거의 안 떨어뜨리나”를 원리로 설명할 수 있다. 이 이해가 LLM을 막연히 쓰는 사람과 의도를 갖고 다루는 엔지니어를 가른다.

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