한 번 한 계산은 다시 하지 않기.
지난 편
4편에서 마지막 벡터로 어휘 전체의 로짓을 내고, 그것을 확률로 바꿔 다음 토큰 하나를 뽑는 것을 봤다.
한 줄 결론
처음 한 번은 입력 전체를 계산하고(prefill), 그 뒤로는 새 토큰 하나만 계산한다(decode). 앞 토큰의 계산 결과는 KV 캐시에 저장해 재사용한다.
쉽게 보기
1편의 반복을 그대로 따르면 같은 계산을 여러 번 한다. 네를 고른 뒤에는 입력 16개에 네를 붙인 17개를 처음부터 다시 계산하고, 쉼표를 고른 뒤에는 18개를 또 계산한다. 답이 길수록 같은 앞부분을 계산하는 횟수도 늘어난다.
그런데 3편에서 본 것처럼 앞 토큰의 벡터는 뒤에 무엇이 붙어도 바뀌지 않는다. 위치 9 토마토는은 몇 번째 반복이든 위치 0~9만 섞으므로 결과가 매번 같다. 그래서 한 번 계산한 결과를 저장해 두면 다음 반복에서는 새로 붙은 토큰만 계산하면 된다. 저장하는 것은 레이어마다 각 위치에서 만든 벡터 두 개이고, 이 저장 공간을 KV 캐시라고 한다(두 벡터는 더 깊이 참고).
2편의 예시 입력으로 보면 이렇다.
| 단계 | 이번에 계산하는 위치 | 고른 토큰 | 계산 뒤 캐시 토큰 수 |
|---|---|---|---|
| prefill | 0~15 (입력 16개 전부) | 네 | 16 |
| decode 1 | 16 네 | , | 17 |
| decode 2 | 17 , | 과일입니다. | 18 |
| decode 3 | 18 과일입니다. | <EOS> | 19 |
prefill. 요청이 들어오면 입력 16개를 한꺼번에 계산한다. 이 16개는 처음부터 다 정해져 있어서 동시에 계산할 수 있다. 한꺼번에 계산해도 위치마다 자기 앞만 섞는 것은 3편에서 본 그대로다. 이때 레이어마다 위치 0~15의 값을 KV 캐시에 넣고, 위치 15의 마지막 벡터로 네를 고른다.
decode. 그다음부터는 방금 고른 토큰 하나만 넣어 계산한다. 위치 16의 네는 레이어마다 위치 0~16을 섞는데, 0~15의 값은 캐시에서 꺼내 쓰고 자기 값만 새로 만들어 캐시에 추가한다. 그래서 decode를 한 번 할 때마다 캐시가 한 토큰씩 늘어난다. <EOS>를 고르면 반복이 멈추고, 2편에서 본 것처럼 <EOS>는 입력에 다시 붙지 않으므로 캐시는 19개에서 끝난다.
이 구조를 알면 API를 쓰면서 보던 현상이 설명된다.
- 입력이 길면 첫 토큰이 늦게 나오는 편이다. 첫 토큰을 고르려면 prefill이 끝나야 하고, prefill은 입력 전체를 계산하기 때문이다. 스트리밍을 켜도 이 대기 시간은 줄지 않는다.
- 첫 토큰 뒤로는 decode가 토큰을 하나씩 내놓는다. 출력이 길수록 decode를 더 많이 해서 끝나는 시점이 늦어진다. 1편에서 답이 길수록 늦게 끝난다고 한 이유가 이것이다.
- KV 캐시는 토큰이 늘수록 커진다. 토큰마다, 레이어마다 값이 쌓이기 때문이다. vLLM 같은 일반적인 서빙 엔진은 KV 캐시를 GPU 메모리에 둔다. 그래서 한 번에 처리할 수 있는 요청 수의 상한은 KV 캐시를 둘 GPU 메모리 크기로 정해진다(KV 캐시의 실제 크기는 더 깊이 참고).
- 여기서 말한 KV 캐시는 요청 하나를 처리하는 동안에만 쓰인다. 2편에서 본 것처럼 두 번째 질문은 앞 대화까지 담은 새 요청이라, 기본적으로는 그 입력 전체를 다시 prefill한다. 요청 사이에서 같은 앞부분의 계산을 다시 쓰는 방법은 7편에서 다룬다.
새 용어
- prefill: 요청을 받아 입력 토큰 전체를 한꺼번에 계산하는 첫 단계. 계산하면서 KV 캐시를 채우고, 끝나면 첫 출력 토큰을 고른다.
- decode: prefill 뒤에 방금 고른 토큰 하나만 계산해 그다음 토큰을 고르는 단계.
<EOS>가 나오거나max_tokens에 도달할 때까지 되풀이한다. - KV 캐시: 앞 토큰을 계산하며 만든 값 중 다음 토큰 계산에 다시 쓰는 값을 레이어마다 저장해 두는 공간. 요청 하나를 처리하는 동안 토큰이 하나 늘 때마다 함께 늘어난다.
더 깊이
K와 V만 저장하는 이유.
3편 더 깊이에서 본 어텐션은 위치 i의 Q를 위치 0~i의 K와 비교해 비율을 내고, 그 비율대로 V를 더한다. decode 2에서 위치 17 ,를 계산할 때 필요한 것을 나눠 보면 이렇다.
- Q: 위치 17의 것 하나만 필요하다. 앞 위치의 Q는 그 위치를 계산할 때 한 번 쓰고 끝난다. 뒤 위치는 자기 Q로 비교하므로 앞 위치의 Q를 다시 쓰지 않는다.
- K, V: 위치 0~17이 모두 필요하다. 이 중 0~16은 이전 단계에서 만든 것과 똑같다. causal mask 때문에 위치 j의 K와 V는 위치 0~j만으로 정해져서, 뒤에 토큰이 붙어도 바뀌지 않는다.
그래서 KV 캐시에는 레이어마다 각 위치의 K와 V만 넣는다. 앞 위치가 레이어를 지나며 만든 벡터도 저장하지 않는다. 뒤 토큰을 계산할 때 이 벡터는 다음 레이어에서 같은 위치의 K, V를 만드는 데만 쓰이는데, 그 K, V는 이미 캐시에 있기 때문이다. RoPE를 쓰는 모델은 위치에 맞춰 회전시킨 뒤의 K를 저장한다. 회전 각도는 그 토큰의 위치만으로 정해지므로 저장한 뒤에 바뀌지 않는다. Hugging Face Transformers의 Llama 구현도 Q와 K를 회전시킨 다음에 K, V를 캐시에 추가한다.
레이어 하나에서 decode를 한 번 하는 과정을 의사 코드로 쓰면 이렇다. x는 이 레이어에 들어온 새 토큰의 벡터이고, kCache와 vCache는 레이어마다 따로 있다. attention은 3편 의사 코드에서 i를 마지막 위치로 둔 계산이다.
const q = toQ(x);
kCache.push(toK(x));
vCache.push(toV(x));
const out = attention(q, kCache, vCache);
메모리 계산: 토큰 하나가 캐시에 남기는 크기.
토큰 하나가 KV 캐시에서 차지하는 크기는 이렇게 계산한다.
2(K와 V) x 레이어 수 x KV 헤드 수 x 헤드 차원 x 2바이트
헤드는 3편 더 깊이에서 본 Q, K, V 묶음이고, 헤드 차원은 묶음 하나의 칸 수다. 2바이트는 숫자 하나를 FP16이나 BF16(16비트 부동소수점)으로 저장한다고 본 값이다.
The Llama 3 Herd of Models (arXiv 2407.21783) Table 3에 따르면 공개 모델 Llama 3 70B는 레이어 80개, 모델 차원 8,192, 어텐션 헤드 64개, KV 헤드 8개다. 헤드 차원은 모델 차원을 어텐션 헤드 수로 나눈 128이다. Hugging Face Transformers의 Llama 설정도 head_dim을 따로 주지 않으면 hidden_size // num_attention_heads로 계산한다.
2 x 80 x 8 x 128 x 2바이트 = 327,680바이트 = 320KiB (토큰 하나)
327,680바이트 x 8,192토큰 = 2,684,354,560바이트 = 2.5GiB
1KiB는 1,024바이트, 1GiB는 1,024 x 1,024 x 1,024바이트다. 입력과 출력을 합쳐 8,192토큰인 요청 하나가 KV 캐시로만 정확히 2.5GiB를 쓴다. 이 크기는 동시에 처리하는 요청마다 따로 든다. vLLM 논문 Efficient Memory Management for Large Language Model Serving with PagedAttention도 3절에서 한 번에 묶어 처리할 수 있는 요청 수의 상한이 KV 캐시를 담을 GPU 메모리로 정해진다고 설명한다.
GQA: KV 헤드를 Q 헤드보다 적게 두는 이유.
3편 더 깊이에서 K와 V 헤드를 Q 헤드보다 적게 두고 여러 Q 헤드가 나눠 쓰는 모델도 있다고 했다. 이 방식을 GQA(grouped-query attention)라고 한다. Llama 3 70B는 Q 헤드 64개가 KV 헤드 8개를 나눠 쓴다. 3편에서 본 Qwen2.5-7B도 Q 헤드 28개에 K, V 헤드 4개라 같은 방식이다.
KV 헤드를 Q 헤드와 같은 64개로 두면 위 계산의 8이 64가 되고, 같은 8,192토큰의 KV 캐시가 20GiB로 8배가 된다. Llama 3 논문은 GQA를 쓴 이유로 추론 속도를 높이고 decode 중 KV 캐시 크기를 줄이는 것을 들었다. GQA를 제안한 Ainslie et al., 2023은 decode 단계마다 가중치와 모든 K, V를 메모리에서 읽는 부담이 decode를 느리게 하는 원인이라고 본다. KV 헤드를 하나만 두면 K, V를 읽는 부담이 크게 줄지만 품질이 떨어질 수 있다. 이 논문은 KV 헤드 수를 1개와 Q 헤드 수 사이로 두는 방식을 시험했다. 품질은 KV 헤드를 Q 헤드와 같은 수로 둔 방식에 가깝고, 속도는 KV 헤드 하나짜리와 비슷하다고 보고했다.
입력이 길면 prefill은 연산이, decode는 메모리 대역폭이 병목이다.
GPU는 계산 장치가 아주 많아서 숫자를 한꺼번에 많이 곱할 때 효율이 좋다. prefill은 입력 토큰을 동시에 계산하므로, 메모리에서 한 번 읽은 가중치 행렬을 입력 토큰마다 한 번씩 쓴다. 입력이 길수록 한 번 읽은 데이터로 하는 계산이 늘어나서, 보통 수백 토큰 이상이면 걸리는 시간이 메모리 읽기보다 계산량에 따라 정해진다. 이 글의 예시처럼 입력이 16개뿐이면 prefill에서도 가중치를 읽는 시간이 계산하는 시간보다 길다.
decode는 토큰 하나를 계산하려고 단계마다 모델 가중치 전체와 그 요청의 KV 캐시를 메모리에서 읽는다. 읽은 양에 비해 계산이 적어서 계산 장치가 쉬는 시간이 생기고, 걸리는 시간은 메모리에서 데이터를 읽어 오는 속도(메모리 대역폭)에 달려 있다. PagedAttention 논문 2.2절은 prefill을 prompt phase라고 부른다. 이 논문에 따르면 prompt phase는 행렬끼리 곱하는 계산이라 GPU의 병렬성을 잘 쓰지만, 생성 단계(이 글의 decode)는 GPU 연산 능력을 대부분 쓰지 못하고 memory-bound가 된다. 위 GQA 논문 서론도 decode가 느린 원인으로 같은 메모리 대역폭 부담을 꼽는다.
그래서 vLLM 같은 일반적인 서빙 엔진은 여러 요청을 묶어 한 번에 계산한다. 가중치를 한 번 읽어 여러 요청에 쓰면 읽는 비용이 요청 수만큼 나뉘기 때문이다(PagedAttention 논문 2.3절). 4편 더 깊이에서 본 배치 크기가 이 묶음의 크기다.
이 차이는 API에서 재는 지연 시간으로도 보인다. Claude 문서는 요청을 보낸 때부터 응답의 첫 토큰이 나올 때까지의 시간을 TTFT(time to first token)라고 부른다. TTFT에는 네트워크와 서버 대기 시간도 들어가지만, 서버 계산 가운데 첫 토큰 전에 끝나야 하는 것은 prefill이다. 그 뒤 토큰이 나오는 속도는 decode가 정한다. 같은 문서는 지연을 줄이는 방법으로 입력과 출력 토큰을 적게 쓰라고 하고, 체감 지연은 스트리밍으로 줄이라고 권한다. 실제 시간은 모델, 하드웨어, 서버 부하에 따라 달라서 여기서는 수치를 들지 않는다.