가중치는 요청 수와 상관없이 한 벌이지만 KV 캐시는 요청마다, 토큰마다 쌓인다. 배치가 크고 문맥이 길면 디코드 스텝에서 읽는 바이트의 대부분이 KV 캐시가 되고, 동시 처리 요청 수의 상한도 KV 캐시 메모리가 정한다. 그래서 KV 캐시를 적은 비트로 저장하면 메모리와 대역폭 양쪽에서 이득을 본다.
토큰당 크기가 절반이 되면
Llama 3.1 8B의 KV 캐시는 BF16 기준 토큰당 128KiB다. FP8로 저장하면 64KiB가 된다.
| KV 형식 | 토큰당 크기 | 55GB에 담기는 토큰 | 4,096토큰 요청의 동시 처리 상한 |
|---|---|---|---|
| BF16 | 128KiB | 약 42만 | 약 102개 |
| FP8 | 64KiB | 약 84만 | 약 205개 |
디코드 스텝에서 읽는 KV 캐시도 절반이 된다. 요청 64개가 각각 2,000토큰 문맥을 가진 배치에서 KV 캐시 읽기는 약 16.8GB에서 약 8.4GB로 줄고, 가중치 16GB와 합친 스텝 하한은 약 9.8ms에서 약 7.3ms로 줄어든다. 문맥이 길수록 줄어드는 폭이 커진다.
가중치 양자화와 달리 KV 캐시 양자화는 모델 체크포인트를 바꾸지 않는다. 엔진이 KV 값을 캐시에 쓸 때 저비트로 바꾸고 어텐션 계산에서 읽을 때 다시 푼다. 그래서 같은 모델에 켜고 끄며 비교하기 쉽다.
스케일
FP8의 E4M3 형식은 최대 448까지만 표현한다. 키와 값을 그대로 넣으면 범위를 넘거나 작은 값들이 0에 뭉칠 수 있으므로, 스케일로 나눠 범위를 맞춘 뒤 저장하고 읽을 때 다시 곱한다.
스케일을 어느 단위로 두느냐가 정밀도와 비용을 정한다.
| 스케일 단위 | 정밀도 | 비용 |
|---|---|---|
| 층별 텐서 하나 | 낮다. 이상값 하나가 전체 스케일을 키운다 | 거의 없다. 미리 보정한 값을 쓸 수 있다 |
| 헤드별 | 중간 | 헤드마다 스케일 하나 |
| 토큰별 또는 채널별 | 높다 | 스케일 저장과 곱셈 비용이 커진다 |
vLLM은 --kv-cache-dtype으로 KV 캐시 형식을 정하고, 문서에 따르면 fp8, fp8_e4m3, fp8_e5m2를 고를 수 있다. 스케일을 보정 데이터로 미리 구해 체크포인트에 넣어 둘 수도 있고, 없으면 기본 스케일을 쓴다. 기본 스케일로 품질이 떨어지면 보정한 스케일부터 시도한다.
키와 값의 분포
KIVI(ICML 2024)는 키 캐시에는 특정 채널에 큰 값이 몰리는 경향이 있어 채널별로 양자화해야 하고, 값 캐시는 토큰별로 양자화해야 한다는 관찰을 바탕으로 KV 캐시를 2비트까지 줄인다.
이 관찰은 8비트에서도 쓸모가 있다. 키에 큰 값이 몰리는 채널이 있는데 텐서 하나에 스케일 하나를 두면, 그 채널의 큰 값이 스케일을 결정해 나머지 채널의 정밀도를 깎는다. 품질 손실이 보이면 스케일 단위를 먼저 의심한다.
2~4비트 KV 양자화는 메모리를 더 줄이지만 엔진과 커널 지원이 제한적이고, 토큰별 스케일 계산과 최근 토큰을 원래 정밀도로 남겨 두는 처리 같은 부가 구조가 필요하다. 운영 환경에서는 엔진이 공식으로 지원하는 형식부터 쓴다.
품질 확인
KV 캐시의 오차는 문맥이 길수록 쌓인다. 새 토큰의 어텐션이 앞선 모든 토큰의 양자화된 키·값을 읽으므로, 짧은 프롬프트에서 멀쩡하던 설정이 긴 문맥에서 무너질 수 있다.
- 긴 문맥 작업을 평가에 넣는다. 긴 문서에서 특정 사실을 찾는 작업, 문서 앞부분과 뒷부분을 함께 참조하는 질문을 서비스의 최대 문맥 길이 근처까지 늘려 본다
- 같은 요청으로 원본과 비교한다. 결정적 디코딩(온도 0)으로 원본 KV와 FP8 KV의 출력을 나란히 놓고 어디서 갈라지는지 본다. 첫 차이가 늦게 나타날수록 오차가 작다
- 가중치 양자화와 따로 잰다. 둘을 함께 켜면 손실의 원인을 가를 수 없다. 하나씩 켜서 품질과 굿풋을 기록한다
켜는 시점
KV 캐시 양자화는 문맥이 긴 서비스일수록 이득이 크고 위험도 크다. 다음 신호가 보이면 가장 먼저 검토한다.
- 선점이 잦다. 동시 처리 요청 수가 KV 메모리에 막혀 진행 중인 요청을 내보내고 있다
- KV 캐시 사용률이 늘 높고 대기열이 생긴다. 연산과 대역폭에는 여유가 있는데 메모리가 먼저 찬다
- 문맥이 길어 디코드 스텝이 KV 읽기에 묶인다. 배치를 키워도 처리량이 오르지 않는다
반대로 프롬프트와 출력이 짧아 KV 캐시가 메모리를 조금만 쓴다면 얻는 것이 적다. 이때는 가중치 양자화나 배치 설정이 더 큰 차이를 만든다.