추론 최적화: KV 캐시 · MQA/GQA · 배칭 · 추측 디코딩 · 서빙 지표

추론 최적화: KV 캐시 · MQA/GQA · 배칭 · 추측 디코딩 · 서빙 지표

학습은 한 번이지만 추론(서빙)은 매일 수백만 번 일어납니다. LLM 추론이 느리고 비싼 가장 큰 이유는 계산이 아니라 메모리 입니다. 토큰 하나를 만들 때마다 수십 GB 가중치를 읽어야 하고, 대화가 길어질수록 KV 캐시 가 GPU 메모리를 채웁니다. 이 수업에서는 프리필과 디코드의 차이, KV 캐시 크기 공식, 이를 줄이는 MQA · GQA, GPU를 쉬지 않게 하는 연속 배칭 · PagedAttention, 작은 모델로 초안을 쓰는 추측 디코딩, 그리고 서빙 성능을 재는 지표를 숫자로 정리합니다.

들어가기 전에

• 트랜스포머 블록 한 층을 손으로 — 쿼리 · 키 · 값과 인과 마스크.
• 수치 형식 · 혼합 정밀도 · 양자화 — "가중치를 읽는 바이트가 병목"이라는 감각.
• 분산 학습 — 큰 모델을 여러 GPU에 나누는 이유(추론도 같은 방법을 씁니다).

한눈에 보기

• 프리필은 계산이, 디코드는 메모리 대역폭이 병목인 이유를 숫자로 설명할 수 있습니다.
• KV 캐시 크기를 공식으로 계산하고, GQA가 이를 몇 배 줄이는지 압니다.
• 연속 배칭 · PagedAttention이 GPU 이용률을 높이는 원리를 압니다.
• 추측 디코딩의 기대 토큰 수 공식으로 속도 이득을 어림합니다.
• TTFT · TPOT · 처리량이 각각 무엇을 재는지, 왜 서로 충돌하는지 압니다.

1. 프리필과 디코드

직관(해설용 비유). 시험 문제를 받아 한 번에 쭉 읽는 단계(프리필)와, 답을 한 글자씩 쓰는 단계(디코드)가 있습니다. 읽기는 여러 줄을 한꺼번에 처리하지만, 쓰기는 앞 글자를 써야 다음 글자를 쓸 수 있어 한 번에 하나뿐입니다.

1. 프리필(prefill): 프롬프트 토큰 전체를 한 번에 병렬로 통과시켜 첫 출력 토큰을 만들고, 모든 층의 키 · 값을 저장합니다.
2. 디코드(decode): 토큰을 하나씩 만듭니다. 매번 새 토큰 하나의 쿼리 · 키 · 값만 계산하고, 과거 토큰의 키 · 값은 저장해 둔 것(KV 캐시)을 씁니다.

왜 디코드는 메모리 병목인가. 가중치 하나(FP16, 2바이트)를 읽으면 토큰 하나에 대해 곱셈 · 덧셈 2 FLOP을 합니다. 배치 1이면 읽은 바이트당 약 1 FLOP 입니다. A100은 BF16 텐서 코어 312 TFLOPS, 메모리 대역폭 약 2.0 TB/s(80GB SXM 기준, NVIDIA 제품 페이지)라 바이트당 약 153 FLOP을 처리할 힘이 있습니다. 디코드는 그 1% 정도만 쓰고 나머지 시간은 메모리를 기다립니다. AWQ 논문도 생성 단계가 메모리 병목이고 연산 강도(arithmetic intensity)가 낮다고 분석합니다.

숫자로 해 보기(해설용 어림). 7B 모델 FP16 가중치 14GB를 2.039 TB/s로 한 번 읽는 데 약 6.9ms. 배치 1의 토큰당 지연 시간은 이보다 짧아질 수 없습니다(이론 하한). 4비트 양자화로 3.5GB가 되면 하한도 약 1.7ms로 줄어듭니다. 이것이 양자화가 추론을 빠르게 하는 주된 이유입니다.

2. KV 캐시: 무엇이고 얼마나 큰가

왜 저장하나. 디코드의 매 단계에서 새 쿼리는 과거 모든 토큰의 키 · 값 과 어텐션을 합니다. 과거 토큰의 키 · 값은 바뀌지 않으므로 매번 다시 계산하지 않고 저장해 둡니다. 계산을 메모리로 바꾸는 거래입니다.

크기 공식.

숫자로 해 보기(LLaMA 65B 구성으로 해설용 계산). LLaMA 논문 표의 65B는 층 80개, 차원 8,192, 헤드 64개(헤드당 128)입니다. 토큰 하나의 KV 캐시는 $2\times80\times8192\times2$ 바이트 $\approx2.6$ MB. 2,048토큰이면 요청 하나에 약 5.4GB, 동시 요청 32개면 약 172GB로 가중치(약 130GB)보다 커집니다. 긴 문맥과 많은 동시 사용자를 감당하기 어려운 이유입니다.

3. MQA와 GQA: 키 · 값 헤드를 줄이기

MQA(멀티쿼리 어텐션, Shazeer 2019). 쿼리 헤드는 여러 개 두되, 키 · 값 헤드는 하나만 두고 모든 헤드가 공유합니다. 논문은 이렇게 하면 KV 텐서가 크게 줄어 점진적 디코딩의 메모리 대역폭 요구가 줄고, 디코딩이 훨씬 빨라지며 품질 저하는 작다고 보고합니다.

GQA(그룹 쿼리 어텐션, Ainslie 외 2023). MHA와 MQA의 중간입니다. 쿼리 헤드를 몇 개 그룹 으로 묶고, 그룹마다 키 · 값 헤드를 하나씩 둡니다. 논문은 기존 MHA 체크포인트를 원래 사전학습 연산의 5%로 추가 학습(uptraining)해 MQA · GQA로 바꿀 수 있고, GQA는 MHA에 가까운 품질을 MQA에 가까운 속도로 낸다고 보고합니다. Llama 2는 34B · 70B 모델에 GQA를 썼습니다.

(8그룹은 해설용 가정입니다. 실제 그룹 수는 모델 설정 파일에서 확인합니다.)

4. 배칭: GPU를 쉬지 않게

왜 배칭인가. 디코드는 가중치를 한 번 읽어 여러 요청에 함께 쓰면 바이트당 연산이 배치 수만큼 늘어납니다. 메모리 병목을 푸는 가장 직접적인 방법입니다.

정적 배칭의 문제. 요청 여러 개를 묶어 시작하면, 답이 짧은 요청이 끝나도 가장 긴 요청이 끝날 때까지 자리가 비어 있고 새 요청은 기다립니다.

연속 배칭(Orca, OSDI '22). Yu 외는 요청 단위가 아니라 반복(iteration) 단위 로 스케줄링해, 매 디코드 단계마다 끝난 요청을 빼고 새 요청을 넣습니다. 길이가 다른 요청을 함께 처리하기 위해 일부 연산만 묶는 '선택적 배칭'도 제안했습니다. 논문은 GPT-3 175B에서 같은 지연 수준으로 NVIDIA FasterTransformer보다 처리량이 36.9배 높았다고 보고합니다.

PagedAttention(vLLM, Kwon 외 2023). KV 캐시를 요청마다 최대 길이만큼 미리 연속 공간으로 잡으면 많이 낭비됩니다. 논문의 측정으로는 기존 시스템에서 KV 캐시 메모리의 20.4%~38.2%만 실제 토큰 상태 저장에 쓰였습니다. vLLM은 운영체제의 가상 메모리 · 페이징처럼 KV 캐시를 작은 블록으로 나눠 필요할 때만 할당하고, 같은 프롬프트를 쓰는 요청끼리 블록을 공유합니다. 낭비를 거의 0으로 줄여 같은 지연에서 처리량을 2~4배 높였다고 보고합니다. 자세한 내용은 PagedAttention 해설 을 보세요.

5. 추측 디코딩: 작은 모델이 초안, 큰 모델이 검토

직관. 신입이 문장 초안을 몇 단어 빠르게 쓰면, 선임이 한 번에 훑어 맞는 데까지 승인하고 틀린 곳부터 고칩니다. 선임이 한 단어씩 직접 쓰는 것보다 빠릅니다.

단계(Leviathan 외, 2022).

1. 작은 근사 모델 $M_q$ 가 토큰 $\gamma$ 개를 차례로 제안합니다.
2. 큰 목표 모델 $M_p$ 가 그 $\gamma$ 개 위치를 한 번에 병렬로 평가합니다(프리필처럼).
3. 앞에서부터 확률 비교로 수락 여부를 정하고, 처음 거절된 자리에서 목표 모델 분포로 다시 뽑습니다.
4. 이 수락 규칙 덕분에 출력 분포는 목표 모델 단독과 똑같습니다.

기대 토큰 수. 수락률을 $\alpha$ 라 할 때(각 토큰이 독립이라는 단순화), 목표 모델 한 번 실행당 생성되는 토큰 수의 기댓값은 논문의 식 (1)입니다.

숫자로 해 보기(해설용). $\alpha=0.8$, $\gamma=3$ 이면 약 2.95개, $\gamma=5$ 면 약 3.69개. $\alpha=0.6$, $\gamma=3$ 이면 약 2.18개입니다. 초안 모델의 비용이 작을수록, 두 모델이 비슷할수록 이득이 큽니다. 논문은 T5-XXL에서 출력이 같으면서 2~3배 빨라졌다고 보고합니다. 디코드가 메모리 병목이라 큰 모델이 토큰 여러 개를 한 번에 검토해도 시간이 거의 늘지 않는다는 점을 이용합니다. 자세한 내용은 추측 디코딩 해설.

6. FlashAttention은 어디에 들어가나

FlashAttention(Dao 외, 2022)은 어텐션 계산을 타일로 나눠 GPU의 느린 HBM과 빠른 온칩 SRAM 사이 읽기 · 쓰기를 줄인 정확한(근사 아닌) 어텐션입니다. 논문은 GPT-2(길이 1K) 학습에서 3배 빨라졌다고 보고합니다. 주로 길이 $n^2$ 의 점수 행렬을 다루는 학습 · 프리필에서 효과가 큽니다. 자세한 내용은 플래시어텐션 · 플래시어텐션-2.

7. 서빙 지표

충돌 관계. 배치를 키우면 처리량은 오르지만 요청마다 TPOT이 늘어납니다. 채팅은 TTFT · TPOT이, 대량 문서 처리는 처리량이 중요합니다. 서비스 목표(예: "TPOT 50ms 이하")를 먼저 정하고 그 안에서 처리량을 최대로 하는 것이 일반적인 접근입니다.

숫자로 해 보기(해설용). TTFT 0.4초, TPOT 30ms, 출력 200토큰이면 전체 지연은 약 $0.4+0.03\times199\approx6.4$ 초입니다.

123456789flowchart LR
  R[요청 도착] --> Q[대기열 · 연속 배칭 스케줄러]
  Q --> P[프리필: 프롬프트 전체 병렬 계산<br/>KV 캐시 기록]
  P --> T1[첫 토큰 → TTFT]
  T1 --> D[디코드: 토큰 하나씩<br/>KV 캐시 읽기 · 추가]
  D -->|추측 디코딩 시 초안 검토| D
  D --> E{끝 토큰?}
  E -->|아니오| D
  E -->|예| F[응답 완료 · KV 블록 반납]

연습 문제

1. 층 32개, KV 헤드 32개, 헤드 차원 128, FP16 모델의 토큰당 KV 캐시는? 4,096토큰이면?
   • 풀이: $2\times32\times32\times128\times2=524{,}288$ 바이트(약 0.52MB). 4,096토큰이면 약 2.1GB입니다.
2. 1번 모델을 KV 헤드 8개인 GQA로 바꾸면?
   • 풀이: 4분의 1인 약 0.54GB입니다.
3. 추측 디코딩에서 $\alpha=0.9$, $\gamma=4$ 일 때 기대 토큰 수는?
   • 풀이: $(1-0.9^5)/(1-0.9)=(1-0.59049)/0.1\approx4.10$ 개입니다.
4. 13B 모델(FP16, 26GB)을 대역폭 2 TB/s GPU에서 배치 1로 돌릴 때 토큰당 지연의 이론 하한은?
   • 풀이: $26\div2000\approx0.013$ 초, 약 13ms입니다.

스스로 점검

1. 디코드에서 배치를 키우면 왜 처리량이 거의 공짜로 늘어나나요?
   • 답: 병목인 가중치 읽기는 한 번이고, 그 가중치를 여러 요청에 함께 쓰므로 바이트당 연산만 늘기 때문입니다. 단, KV 캐시가 메모리를 채우는 데까지입니다.
2. 추측 디코딩이 품질을 해치지 않는 이유는?
   • 답: 수락 · 재샘플링 규칙이 최종 출력 분포를 목표 모델 단독 분포와 똑같이 만들기 때문입니다.
3. PagedAttention이 처리량을 높이는 원리는?
   • 답: KV 캐시를 블록 단위로 필요한 만큼만 할당해 낭비를 없애므로, 같은 메모리에 더 많은 요청을 동시에 담아 배치를 키울 수 있기 때문입니다.

용어 정리

더 깊이 보려면

• 플래시어텐션: 메모리를 아끼는 정확한 어텐션
• 추측 디코딩으로 트랜스포머 추론 빠르게 하기
• PagedAttention: LLM 서빙의 메모리 관리 (vLLM)
• vLLM 문서: docs.vllm.ai

다음에 볼 것

• 플래시어텐션: 메모리를 아끼는 정확한 어텐션 — 어텐션 계산 자체의 메모리 이동 줄이기
• 직접 해 보기: 양자화 오차와 KV 캐시 메모리 계산하기 — 오늘 공식을 코드 계산기로

라이선스

본문은 KAITABASE 자체 해설입니다. 기법별 수치는 MQA · GQA · Orca · PagedAttention · 추측 디코딩 · FlashAttention · AWQ 원문과 LLaMA · Llama 2 논문, NVIDIA A100 제품 페이지에서 확인한 것만 옮겼습니다. KV 캐시 · 지연 시간 계산은 해설용이며 파이썬으로 검산했습니다.

출처

• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (2019): arxiv.org/abs/1911.02150
• Ainslie 외, GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023): arxiv.org/abs/2305.13245
• Yu 외, Orca: A Distributed Serving System for Transformer-Based Generative Models (OSDI '22): usenix.org/conference/osdi22/presentation…
• Kwon 외, Efficient Memory Management for Large Language Model Serving with PagedAttention (2023): arxiv.org/abs/2309.06180
• Leviathan 외, Fast Inference from Transformers via Speculative Decoding (2022): arxiv.org/abs/2211.17192
• Dao 외, FlashAttention (2022): arxiv.org/abs/2205.14135
• Lin 외, AWQ (2023): arxiv.org/abs/2306.00978
• Touvron 외, LLaMA (2023): arxiv.org/abs/2302.13971 · Llama 2 (2023): arxiv.org/abs/2307.09288
• NVIDIA A100 제품 페이지: nvidia.com/en-us/data-center/a100