LLM에서 KV 캐시는 얼마나 생길까
KV 캐시는 사용자마다 따로 생길까? 질문 하나만 던져도 생길까?
LLM 서빙에 필요한 메모리는 가중치만으로 가늠할 수 없다. 가중치로는 모델을 올리는 데 필요한 메모리만을 계산할 수 있고, 그와 별개로 연산 과정에서 사용될 KV 캐시를 위해 여유 메모리를 반드시 확보해야한다. KV 캐시는 요청 하나마다 따로 생기고 그 요청의 토큰 수에 비례해 커진다. 물론 동시 사용자가 늘면 캐시도 그 수만큼 늘어난다. 그래서 실제 서빙에서는 GPU 메모리 대부분을 가중치가 아니라 KV 캐시가 차지하는 일이 흔하다.
▍KV 캐시란
트랜스포머는 토큰을 하나 생성할 때마다 앞선 모든 토큰을 다시 확인한다. 앞선 토큰 하나하나가 새로 생성될 토큰과 얼마나 관련이 있는지 점수를 매기고, 관련이 큰 토큰의 정보를 더 많이 섞어서 다음 토큰을 예측한다. 여기서 점수를 매길 때 쓰는 벡터를 키(Key), 정보를 담은 벡터를 값(Value) 이라고 부르는데, 토큰마다 키와 값을 하나씩 가진다. 앞선 토큰의 키와 값은 한 번 계산하면 바뀌지 않으니 저장해 두고 다시 쓰는데, 이 저장소가 KV 캐시다.
쿼리(q)·키(k)·값(v)은 모두 같은 토큰에서 나온다. 토큰은 레이어에 들어올 때 숫자 4,096개짜리 벡터 x로 표현되는데(Llama 3.1 8B 기준), 이 x에 학습으로 정해진 가중치 행렬 세 개를 각각 곱하면 q, k, v가 나온다. 행렬은 레이어마다 따로 있고 추론 중에는 바뀌지 않는다.
곱한 결과는 128칸씩 잘라 헤드별로 쓴다. q는 헤드 32개로, k와 v는 헤드 8개로 나뉘어 쿼리 헤드 네 개가 KV 헤드 하나를 같이 쓴다. q와 k에는 위치 정보(RoPE)를 한 번 더 입혀서 토큰이 문장의 몇 번째에 있는지도 점수에 반영한다. 지나간 토큰의 x는 뒤에 새 토큰이 와도 바뀌지 않으니 그 토큰의 k와 v도 끝까지 같고, 그래서 저장해 두고 다시 쓸 수 있다. q는 지금 생성하는 토큰에만 필요해서 캐시에 넣지 않는다.
예를 들어 문장이 “나는 커피를 정말 좋아”까지 생성됐고 그다음 토큰을 정하는 스텝이라면, 한 레이어의 한 헤드에서는 이런 계산이 일어난다.
점수와 가중치는 앞선 토큰 하나에 숫자 하나씩이지만, 출력은 스칼라가 아니라 v와 같은 128칸짜리 벡터다. 가중치 0.65는 “커피를”의 v 벡터 128칸 전체에 곱해지고, 네 토큰에서 나온 벡터를 칸끼리 더하면 출력이 된다. 실제 계산에서는 점수를 √128로 나눈 다음 softmax에 넣는데, 그림에서는 이 단계를 생략했다. 헤드 32개의 출력 벡터를 이어 붙이면 다시 4,096칸이 되고, 이 벡터가 행렬 하나와 피드포워드 층을 더 거쳐 다음 레이어의 x가 된다.
요청이 들어오면 엔진은 먼저 프롬프트 전체를 한 번에 계산하면서 모든 토큰의 K·V를 레이어마다 저장한다. 이 단계를 프리필(prefill)이라고 부른다. 그다음 답을 한 토큰씩 생성하는 디코드 단계에서는 토큰이 하나 나올 때마다 그 토큰의 K·V가 캐시 끝에 새롭게 추가된다. 프롬프트가 2,000토큰이고 답이 500토큰이면 응답이 끝나는 순간 캐시에는 2,500토큰에 해당하는 값이 들어 있다.
이 캐시는 그 요청의 토큰으로만 계산되기 때문에 다른 요청과 나눠 쓸 수 없다. 단, 앞부분이 토큰 단위로 완전히 같은 경우는 예외인데, 이건 뒤에서 다루므로 넘어간다.
▍토큰 하나의 크기
토큰 하나가 차지하는 KV 캐시의 크기는 모델 구조만 알면 구할 수 있다. 레이어마다, KV 헤드마다 키 벡터와 값 벡터가 하나씩 저장되니 이 수들을 곱하고 원소 하나의 바이트를 곱하면 된다.
bytes_per_token = 2 (K, V) × layers × kv_heads × head_dim × bytes_per_element
Llama 3.1 8B에 대입하면 2바이트짜리 값이 헤드 차원, K와 V, KV 헤드, 레이어 수만큼 차례로 곱해진다.
Llama 3.1 두 모델의 설정을 넣으면 이렇게 나온다. 두 모델 모두 GQA를 써서 쿼리 헤드는 각각 32개, 64개지만 KV 헤드는 8개뿐이다.
| 모델 | 레이어 | KV 헤드 | 헤드 차원 | 토큰당 (BF16) |
|---|---|---|---|---|
| Llama 3.1 8B | 32 | 8 | 128 | 2 × 32 × 8 × 128 × 2바이트 = 128KiB |
| Llama 3.1 70B | 80 | 8 | 128 | 2 × 80 × 8 × 128 × 2바이트 = 320KiB |
128KiB는 작아 보이지만 토큰 수를 곱하면 금방 커진다. 요청 하나가 쓰는 토큰은 프롬프트와 생성한 답을 합친 길이다.
| 요청 하나의 토큰 수 | 8B | 70B |
|---|---|---|
| 2,048 | 0.25GiB | 0.63GiB |
| 8,192 | 1GiB | 2.5GiB |
| 32,768 | 4GiB | 10GiB |
| 131,072 (두 모델의 최대 컨텍스트) | 16GiB | 40GiB |
8K 토큰짜리 대화 하나에 8B 모델은 1GiB를 쓴다. 128K 컨텍스트를 꽉 채우면 요청 하나가 16GiB를 써서 8B 모델의 BF16 가중치(약 15GiB)보다 커진다. 70B 모델이면 같은 요청이 40GiB를 쓰는데, 가중치 약 131GiB에 이만큼이 더해진다.
▍대화가 이어질 때의 캐시
대화가 길어질수록 그 사람 몫의 캐시가 계속 쌓인다고 생각하기 쉬운데, 이건 한 요청에서만 해당한다. OpenAI 호환 API처럼 상태를 저장하지 않는 채팅 API에서는 턴 하나가 요청 하나다. 서버는 응답을 다 보내면 그 요청의 캐시를 해제한다. 다음 턴에는 클라이언트가 이전 대화 전체를 다시 보내고, 서버는 그걸 처음부터 프리필한다.
8B 모델에 대화를 보낸다고 해 보자. 시스템 프롬프트와 첫 질문이 1,200토큰이고, 이후 질문은 200토큰, 답은 매번 500토큰이다.
턴마다 캐시는 이전 대화와 새 답을 합친 길이만큼 차고, 열 번째 턴이 끝날 때는 8,000토큰, 약 1,000MiB가 된다. 서버는 턴이 끝날 때마다 메모리를 돌려받지만, 이미 한 계산은 다음 턴에 또 해야 한다. 같은 이전 대화를 턴마다 다시 프리필하니 열 턴 동안 프리필한 토큰은 43,500토큰이다. 새로 들어온 입력은 3,000토큰뿐이니 열네 배 넘게 계산했다. 그래서 대화가 길어질수록 첫 토큰이 늦게 나온다.
▍RTX 4090 한 장에서의 계산
8B 모델은 가중치가 약 15GiB라 24GiB 카드에 넉넉히 올라갈 것처럼 보이지만, 이 식으로 계산하면 KV 캐시에 남는 메모리가 생각보다 적다. 아래는 계산값이고, 실제 수치는 엔진이 시작할 때 남기는 로그로 확인한다.
토큰당 128KiB로 나누면 GPU 한 장에 들어가는 KV 캐시는 모든 요청을 합쳐 3만 8천~4만 6천 토큰 정도다. 그런데 모델 설정의 최대 컨텍스트는 131,072토큰이라 최대 길이 요청은 하나도 담을 수 없다. vLLM은 최대 길이 요청 하나가 KV 캐시에 들어가지 않으면 아래 문구로 시작하는 에러를 내고 뜨지 않는다.
The model's max seq len (131072) is larger than the maximum number of tokens that can be stored in KV cache ...
만약 --max-model-len을 8192로 낮추면 엔진이 뜰 것이다. 그래도 4만 토큰 남짓을 요청들이 나눠 쓰므로 8K를 다 쓰는 대화라면 동시에 다섯 개 정도가 한계다. 이 카드에서 동시 사용자를 늘리려면 가중치가 아니라 KV 캐시에 할당되는 메모리를 늘려야 한다. 엔진은 시작할 때 KV 캐시에 들어가는 토큰 수와 최대 동시 처리 수를 로그로 남긴다. vLLM V1이라면 GPU KV cache size와 Maximum concurrency로 시작하는 줄이니, 계산이 맞는지 그 줄과 비교하면 된다.
▍KV 캐시를 줄이는 기법
PagedAttention, 프리픽스 캐시, KV 캐시 FP8 양자화는 모두 요청마다 생기는 캐시를 줄이려고 나왔는데, 효율화 하는 부분이 저마다 다르다.
| 기법 | 줄이는 것 | 효과 |
|---|---|---|
| PagedAttention | 예약만 하고 쓰지 않는 메모리 | 캐시를 고정 크기 블록(vLLM 기본 16토큰)으로 나눠 필요할 때마다 할당한다 |
| 프리픽스 캐시 | 요청 사이의 중복 | 앞부분이 같은 요청끼리 그 부분의 블록을 한 벌만 저장하고 프리필도 건너뛴다 |
| FP8 KV 캐시 | 원소당 바이트 | 2바이트를 1바이트로 줄여 토큰당 128KiB가 64KiB가 된다 |
PagedAttention 등장 이전의 엔진은 요청이 얼마나 길어질지 모르니 최대 길이만큼 연속 공간을 미리 잡아 두었다. 8K까지 허용하는 서버에 500토큰짜리 요청이 오면 나머지 7,500토큰 몫은 빈칸으로 남는데, 다른 요청은 이걸 쓰지 못한다. PagedAttention을 발표한 연구진은 기존 시스템이 KV 캐시 메모리의 60~80%를 이렇게 낭비한다고 측정했다. PagedAttention은 운영체제가 메모리를 페이지 단위로 나눠 관리하는 방식을 KV 캐시에 그대로 가져왔다. 엔진은 시작할 때 KV 캐시에 쓸 메모리를 16토큰짜리 블록으로 미리 잘라 두고, 요청마다 블록 테이블을 하나씩 만든다. 블록 테이블에는 그 요청의 몇 번째 블록이 실제 메모리의 몇 번 블록에 들어 있는지가 적힌다.
요청이 들어오면 엔진은 프롬프트를 담을 만큼만 블록을 내준다. 디코드 중에 마지막 블록이 16토큰으로 다 차면 빈 블록을 하나 더 꺼내 블록 테이블 끝에 적고, 요청이 끝나면 쓰던 블록을 전부 빈 블록 목록으로 돌려놓는다. 어텐션을 계산할 때는 커널이 블록 테이블을 따라가며 블록을 순서대로 읽기 때문에, 한 요청의 블록이 메모리 여기저기에 흩어져 있어도 계산 결과는 같다.
그래서 미리 잡아 두는 공간이 사라지고, 낭비는 요청마다 다 차지 않은 마지막 블록의 빈칸으로 줄어든다. 블록이 16토큰이면 요청 하나당 많아야 15토큰이다.
그림의 블록 할당 줄을 블록 테이블로 적으면 이렇다. 블록 번호는 메모리 왼쪽부터 0번이다.
| 요청 | 블록 테이블 (요청 안의 순서 → 실제 블록 번호) |
|---|---|
| A | 0 → 0, 1 → 4 |
| B | 0 → 1, 1 → 2, 2 → 5, 3 → 6, 4 → 8, 5 → 9 |
| C | 0 → 3, 1 → 7, 2 → 10 |
블록 테이블 여러 개가 같은 실제 블록을 가리키게 할 수도 있다. 엔진은 블록마다 몇 개의 요청이 가리키고 있는지 세어 두었다가 그 수가 0이 될 때 빈 블록 목록으로 돌려놓는다. 아래의 프리픽스 캐시가 이 구조 위에서 동작한다.
프리픽스 캐시를 켜면 턴마다 이전 대화를 다시 프리필하던 계산이 줄어든다. vLLM은 요청이 끝나도 블록을 바로 지우지 않고 남겨 두었다가, 메모리가 모자랄 때 오래 쓰이지 않은 것부터 내보낸다(eviction). 블록을 내보내기 전에 같은 대화의 다음 턴이 같은 서버로 오면, vLLM은 이전 대화의 블록을 그대로 쓰고 새 질문 200토큰 남짓만 프리필한다. 다만 앞부분이 토큰 단위로 정확히 같아야 해서, 시스템 프롬프트 첫 줄에 현재 시각을 넣거나 하면 엔진은 그 뒤 전체를 다시 계산한다. 만약 서버가 여러 대라면 같은 대화를 같은 서버로 보내는 라우팅도 함께 두어야 한다.
엔진은 남겨 둔 블록을 해시로 찾는다. vLLM은 다 찬 블록마다 그 블록의 토큰 16개와 그 앞의 토큰 전체를 묶어 해시값을 계산하고, 해시값과 실제 블록 번호를 해시 테이블에 적어 둔다. 새 요청이 들어오면 프롬프트를 16토큰씩 잘라 앞에서부터 같은 방식으로 해시를 구해 테이블에서 찾는다. 처음으로 못 찾는 블록이 나올 때까지는 기존 블록 번호를 블록 테이블에 그대로 적고, 그 뒤부터 프리필한다. 16토큰을 다 채우지 못한 마지막 조각은 해시를 만들지 않아 재사용하지 않는다.
해시에 앞의 토큰 전체가 들어가기 때문에, 첫 블록이 하나라도 다르면 뒤 블록은 내용이 같아도 해시가 전부 달라진다. 그래서 첫 줄에 현재 시각을 넣으면 엔진은 뒤 블록을 하나도 찾지 못하고 전부 다시 계산한다. 요청 여러 개가 같은 블록을 쓰는 동안에는 앞에서 본 참조 수가 그만큼 올라가고, 엔진이 내보내는 대상은 참조 수가 0인 블록뿐이다. 그래서 다른 요청이 쓰고 있는 시스템 프롬프트 블록은 메모리가 모자라도 지워지지 않는다.
FP8 KV 캐시는 키와 값을 저장하는 숫자 형식을 16비트에서 8비트로 바꾼다. 모델 가중치는 그대로 두고, 엔진이 K·V를 캐시에 쓸 때만 8비트로 바꿔 넣었다가 어텐션을 계산할 때 원래 정밀도로 다시 풀어 쓴다.
| 형식 | 부호 · 지수 · 가수 비트 | 표현할 수 있는 최댓값 |
|---|---|---|
| BF16 | 1 · 8 · 7 | 약 3.4 × 10³⁸ |
| FP8 E4M3 | 1 · 4 · 3 | 448 |
| FP8 E5M2 | 1 · 5 · 2 | 57,344 |
지수 비트는 표현할 수 있는 범위를, 가수 비트는 그 범위 안을 얼마나 촘촘하게 나타낼지를 정한다. E4M3은 가수가 3비트뿐이라 1과 2 사이를 0.125 간격으로, 2와 4 사이를 0.25 간격으로만 나타낸다. 그래서 3.14159를 넣으면 3.25로, 0.1234를 넣으면 0.125로 저장된다. 448보다 큰 값은 아예 담을 수 없어서, 엔진은 값을 스케일로 나눠 범위 안에 맞춘 뒤 저장하고 읽을 때 스케일을 다시 곱한다. 스케일은 보통 레이어마다 하나를 두고, 보정 데이터로 미리 구해 체크포인트에 넣어 둘 수도 있다.
vLLM에서는 --kv-cache-dtype fp8 옵션 하나로 켤 수 있다. 켜면 토큰당 크기가 128KiB에서 64KiB로 줄어 같은 메모리에 토큰이 두 배로 들어간다. 앞의 4090이라면 KV 캐시에 들어가는 토큰이 7만 6천~9만 2천 개로 늘어, 8K 대화를 동시에 9-11개까지 받을 수 있다.
대신 저장할 때마다 반올림 오차가 생기고, 새 토큰은 앞선 모든 토큰의 키와 값을 읽으므로 문맥이 길수록 오차가 쌓인다. 짧은 프롬프트에서 멀쩡하던 설정이 긴 문서에서는 틀린 답을 낼 수 있으니, 적용하기 전에 서비스의 최대 문맥 길이에 가까운 입력으로 평가를 한 번 돌려 보는 것이 좋다.
--max-model-len을 정하기 전에 토큰당 크기에 컨텍스트 길이와 동시 사용자 수를 곱해 본다. 그 값이 가중치를 올리고 남는 메모리보다 크면 컨텍스트 길이와 사용자 수 중 하나를 줄여야 한다.
참고
- Llama 3.1 8B Instruct 모델 카드와 config.json
- Llama 3.1 70B Instruct 모델 카드와 config.json
- Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention (SOSP 2023)
- Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023)
- vLLM: Automatic Prefix Caching
- vLLM: Quantized KV Cache
- Micikevicius et al., FP8 Formats for Deep Learning (2022)
- vLLM: Engine Arguments
계산 기준: BF16 KV 캐시, 1GiB = 2³⁰바이트. 4090 예산은 계산값이고 실측이 아니다.