Transformer Inference의 두 단계: Prefill과 Decode 및 KV Caching 차이

 

 

Transformer 모델에서 inference는 크게 Prefill과 Decode의 두 단계로 이루어진다. 이 두 단계는 입력 토큰을 처리하고, 출력을 생성하는 방식이 다르며, 각각 성능에 중요한 영향을 미친다. 또한, 효율적인 추론을 위해 KV (Key-Value) 캐싱이 활용되는데, 이를 통해 불필요한 연산을 줄이고 속도를 최적화할 수 있다.

---

 

Prefill 단계: 입력 토큰 처리 단계

 

Prefill 단계는 모델이 전체 입력 시퀀스를 한 번에 병렬로 처리하여, 각 토큰을 key-value 쌍으로 변환하는 과정이다. 이때 행렬-행렬 곱셈(matrix-matrix multiplication) 연산을 통해 모든 토큰이 병렬로 처리되므로, 여러 토큰을 동시에 계산할 수 있어 GPU를 최대로 활용할 수 있다. 이로 인해 배치(batch) 단위로 대량의 입력을 빠르게 처리할 수 있으며, 추론 속도 또한 Decode 단계보다 빠르고 효율적이다.

---

 

Decode 단계: 순차적 토큰 생성 단계

 

Decode 단계는 모델이 토큰을 하나씩 생성하는 과정이다. 이 단계에서는 행렬-벡터 곱셈(matrix-vector multiplication)이 사용되며, 각 새 토큰은 이전에 생성된 토큰들에 의존한다. 이는 순차적(autoregressive)인 생성 방식이기 때문에 병렬 처리가 어려워 속도가 느리며 GPU 활용 효율이 떨어진다. 특히 대규모 모델에서는 이로 인해 디코딩 단계에서 큰 지연(latency) 문제가 발생할 수 있다.

---

 

KV (Key-Value) Caching: 불필요한 연산 최소화

 

KV Caching은 이전 토큰의 키(Key)와 값(Value) 텐서를 저장하여, 모델이 매번 새 토큰을 생성할 때 이전 시퀀스 전체를 다시 계산할 필요가 없도록 한다. 이를 통해 연산량이 줄어들고, 추론 속도가 빨라진다. 캐시의 크기는 토큰 수, 레이어 수, 주의 헤드 수에 따라 선형적으로 증가하므로, 특히 Decode 단계에서 효율적으로 관리해야 한다.

 

KV Caching의 단계별 적용 방식

1. Prefill 단계: 이 단계에서 전체 입력 시퀀스에 대한 모든 K와 V 값이 계산되고, 이를 캐시에 저장한다.
2. Decode 단계: 이전에 캐시된 K와 V 값들은 재사용되며, 새로 생성되는 토큰의 쿼리(Q) 값만 추가로 계산하고 K와 V 캐시에 추가(append)한다. 이를 통해 모델이 매번 전체 시퀀스를 다시 계산하지 않도록 최적화할 수 있다.

 

---