DevOps··2분 읽기·0

6. 추론: Prefill, Decode, KV 캐시

LLM이 토큰을 생성할 때 왜 전체를 다시 계산하지 않는지. KV 캐시의 동작 원리, VRAM 비용, Prefill과 Decode의 차이.

글꼴

5편에서 5개 토큰을 한꺼번에 36개 레이어에 통과시켜 "종료시켰다"를 얻었다. 이제 6개 토큰이다. 6개를 통째로 다시 돌리면 앞의 5개를 또 계산하는 셈이다.

이 낭비를 없애는 게 KV 캐시다.


핵심 아이디어: K, V를 저장해두면 된다

3편에서 어텐션은 "새 토큰의 Q x 이전 토큰들의 K -> 주목도 -> 이전 토큰들의 V 가중합"이라고 했다. 이전 토큰들의 K와 V는 이미 계산이 끝났고, 이후에 바뀌지 않는다.

저장해두면 재계산할 필요가 없다. 이걸 KV 캐시라고 한다.


Prefill — 입력을 한꺼번에 처리하면서 캐시를 채운다

사용자 입력: "OOM killer가 nginx를"  (5개 토큰)
 
  5개를 한꺼번에 36개 레이어에 통과 (5편에서 본 과정)
 
  이때 각 레이어에서:
    5개 토큰 각각의 Q, K, V를 계산
    Q -> 어텐션 계산에 쓰고 버림
    K, V -> 캐시에 저장  <-- 이게 핵심
 
  레이어 1의 KV 캐시: [K_OOM, K_killer, K_가, K_nginx, K_를]
                      [V_OOM, V_killer, V_가, V_nginx, V_를]
  ...
  레이어 36의 KV 캐시: [K_OOM, K_killer, K_가, K_nginx, K_를]
                       [V_OOM, V_killer, V_가, V_nginx, V_를]
 
  -> 36개 레이어 x 5개 토큰의 K, V가 VRAM에 저장됨.
  -> 마지막 토큰("를") 위치에서 "종료시켰다" 선택.

Prefill은 입력 전체를 한꺼번에 처리하면서 KV 캐시를 채우는 단계다. 5편에서 본 과정이 바로 이거였다.


Decode — 새 토큰 1개만 통과시킨다

"종료시켰다"를 생성한 뒤에는 이 1개만 처리하면 된다.

"종료시켰다" 1개만 처리하면 된다.
 
  레이어 1:
    "종료시켰다"의 Q, K, V를 계산
    Q x 캐시된 5개의 K -> 주목도 -> 캐시된 V의 가중합
    K, V를 캐시에 추가
    -> FFN -> 잔차연결
 
  레이어 1의 KV 캐시: [...5개...] + [K_종료시켰다]  <- 6개로 증가
 
  레이어 2~36도 동일. -> 최종 벡터 -> 출력 행렬 -> "." 선택.
 
다음 Decode:
  "." 1개만 처리.
  각 레이어에서 캐시된 6개의 K와 내적 -> K, V 추가 -> 캐시 7개로 증가.
 
  ...EOS까지 반복.

Q는 왜 캐시 안 하나? Q는 "이 토큰이 뭘 찾고 있는지"이고, 해당 토큰 처리 순간에만 사용한다. 이후 다른 토큰이 이 토큰의 Q를 참조할 일은 없다. K는 "이 토큰이 뭘 제공하는지"라서 이후 토큰들이 계속 참조한다. 그래서 K, V만 캐시한다.

기존 K, V는 절대 수정되지 않는다. 새 것만 추가될 뿐이다.


KV 캐시의 VRAM 사용량

이 캐시가 VRAM을 얼마나 먹는지 계산해보면 꽤 놀랍다.

토큰 1개의 KV 캐시 (3B 모델):
  K: 2048차원 x 2바이트 = 4KB
  V: 2048차원 x 2바이트 = 4KB
  x 36 레이어 = 288KB
 
토큰 수별:
    100 토큰  ->    28 MB
  1,000 토큰  ->   281 MB
 10,000 토큰  ->  2.8 GB
128,000 토큰  -> 36.0 GB  <- 파라미터(6GB)보다 훨씬 큼

128K 토큰을 처리하면 KV 캐시만 36GB다. 파라미터 크기(6GB)의 6배다. 긴 대화가 VRAM을 잡아먹는 이유가 이거다. ollama에서 대화를 길게 이어가다 보면 점점 느려지는 경험을 해봤을 거다.


Prefill vs Decode — 왜 입력은 빠르고 출력은 느린가

ollama로 프롬프트를 넣으면 처음에 잠깐 멈추다가(Prefill), 그 뒤부터 글자가 하나씩 나온다(Decode). 이 두 단계의 특성이 완전히 다르다.

                Prefill              Decode
              (입력 처리)            (토큰 생성)
---------------------------------------------
처리 단위      N개 토큰 동시          1개 토큰씩
W 읽기        1번 읽고 N개 계산      매 토큰마다 전부 읽기
병목          계산량 (compute)       메모리 대역폭 (memory)
체감          빠름                  글자가 하나씩 나옴
---------------------------------------------

Prefill은 W를 한 번 읽고 N개 토큰을 동시에 계산한다. GPU의 병렬 처리 능력을 최대한 활용할 수 있다. 병목은 순수 계산량이다.

Decode는 다르다. 토큰 1개를 생성할 때마다 모든 W를 VRAM에서 GPU 코어로 읽어야 한다.

토큰 1개 생성할 때 읽어야 하는 W:
 
레이어 1개:
  W_Q  (2048x2048) =  8MB
  W_K  (2048x2048) =  8MB
  W_V  (2048x2048) =  8MB
  W_O  (2048x2048) =  8MB
  W1   (2048x8192) = 32MB
  W2   (8192x2048) = 32MB
  -------------------------
  레이어 1개 합계: ~96MB
 
x 36개 레이어 = ~3.4GB
+ 임베딩 + 출력 행렬 = 0.6GB
-------------------------
총 ~4GB를 매 토큰마다 VRAM에서 읽음

W는 고정이다. 바뀌지 않는다. 하지만 매 토큰마다 다시 읽어야 한다. 토큰마다 입력 벡터가 다르니까 같은 W를 곱해도 결과가 다르기 때문이다. Decode의 병목은 계산 속도가 아니라 메모리 대역폭이다. 이건 9편(하드웨어)에서 더 자세히 다룬다.

여러 사용자를 동시에 처리하려면 — Continuous Batching

지금까지는 사용자 1명의 요청을 처리하는 관점이었다. 서버에서 여러 사용자의 요청을 동시에 처리하려면 어떻게 할까?

단순 배칭:
  사용자A 요청 (100토큰) 시작 -> 완료 -> 사용자B 요청 시작 -> 완료
  -> 한 명이 끝나야 다음 사람. GPU가 놀 때가 많다.
 
Continuous Batching:
  사용자A Prefill -> 사용자B Prefill -> 사용자C Prefill
  사용자A Decode (1토큰) + 사용자B Decode (1토큰) + 사용자C Decode (1토큰)
  사용자A 완료 -> 사용자D Prefill (빈 자리에 즉시 투입)
  ...
 
  핵심: 매 스텝마다 여러 사용자의 토큰을 한꺼번에 처리.
  한 사용자가 끝나면 대기열에서 즉시 새 사용자를 넣음.
  GPU 활용률이 크게 올라감.

Prefill은 토큰 N개를 동시에 처리하니까 GPU를 잘 활용하지만, Decode는 토큰 1개씩이라 GPU가 놀게 된다. Continuous Batching은 여러 사용자의 Decode를 묶어서 한 번에 처리함으로써 GPU 활용률을 높인다.

단, 사용자가 늘어나면 각자의 KV 캐시가 VRAM에 동시에 올라가야 한다. 사용자 10명이 각각 2000토큰을 처리 중이면 KV 캐시만 10 x 0.6GB = 6GB가 필요하다. VRAM이 병목이 되는 순간이다.

vLLM이나 TGI(Text Generation Inference) 같은 서빙 프레임워크가 이 Continuous Batching을 구현하고 있다. ollama는 로컬 단일 사용자 중심이라 이 부분이 단순하지만, 서버에 올려서 여러 사람이 쓰는 상황이라면 이런 서빙 프레임워크를 고려해야 한다.


여기까지가 LLM이 실제로 토큰을 생성하는 과정이다. Prefill로 캐시를 채우고, Decode로 1개씩 생성하고, KV 캐시가 VRAM을 누적으로 잡아먹는다. 그런데 "종료시켰다" 23% 중에서 왜 하필 "종료시켰다"가 선택되는 걸까? 항상 가장 높은 확률을 고르는 건가? 다음 편에서 확률 분포에서 토큰을 고르는 방법을 다룬다.

이 글이 어떠셨나요?

이 글이 도움이 되셨나요?
공유:

관련 포스트

뉴스레터 구독

새 글이 올라오면 이메일로 알려드려요.

댓글

댓글을 불러오는 중...