KV 캐시 메모리 산정 계산기 만들기

목표

KV 캐시 메모리 공식을 코드로 구현하고, GQA 와 양자화의 효과를 계산으로 확인하고, 주어진 GPU 에서 가능한 최대 동시 요청 수를 산출한다.

왜 중요한가

"A100 80GB 한 장에 이 모델이 들어가나요"라는 질문은 배포 전에 반드시 나옵니다. 가중치만 계산하면 답이 틀립니다. 7B 모델의 가중치는 fp16 으로 14GB 지만, 128K 컨텍스트에 배치 8 이면 KV 캐시가 512GiB 입니다. 가중치의 36배입니다. 이 계산을 안 하고 시작하면 배포 당일에 OOM 을 만나고, OOM 은 부하가 몰릴 때 나므로 가장 나쁜 시점에 납니다. 그리고 이 계산기는 한 번 만들어 두면 계속 씁니다 — 모델을 바꿀 때, 컨텍스트 길이를 늘릴 때, 양자화를 검토할 때, max_num_seqs 를 정할 때 매번 필요합니다. 6번 스텝에서 구하는 최대 동시 요청 수가 용량 계획의 출발점이 되는 숫자입니다.

단계

1. /root/kv/kv.py 에 kv_bytes(layers, hidden, seq, batch, dtype_bytes, kv_groups=1) 를 만든다. 기본 공식은 2 * layers * hidden * seq * batch * dtype_bytes 이다.
2. layers=32, hidden=4096, seq=4096, batch=1, dtype_bytes=2 로 계산한 값을 /root/kv/base.txt 에 bytes=<정수> gib=<소수 둘째자리> 로 적는다. gib 는 2.00 이어야 한다.
3. seq=131072 로 바꾼 값을 /root/kv/long.txt 에 같은 형식으로 적는다. gib 는 64.00 이어야 한다.
4. kv_groups=8 을 적용해 /root/kv/gqa.txt 에 적는다. gib 는 8.00 이어야 한다(128K 기준).
5. dtype_bytes=1 을 적용해 /root/kv/int8.txt 에 적는다. gib 는 4.00 이어야 한다(128K, GQA-8 기준).
6. /root/kv/fit.txt 에 gpu_gib=80 weights_gib=14 overhead_gib=6 kv_budget_gib=60 per_req_gib=<소수> max_concurrent=<정수> 를 적는다. 요청당 기준은 32K 컨텍스트, GQA-8, fp16 이다.
7. /root/kv/kv_table.csv 에 context,batch,gib 헤더와 6행을 적는다. 컨텍스트는 4096, 32768, 131072 이고 배치는 1과 8이다. 값은 GQA 없이 fp16 기준이다.

참고

• 공식의 앞 2는 키와 값 두 벌이라는 뜻입니다.
• GQA 는 키·값 헤드를 그룹으로 공유하므로 그룹 수로 나눕니다.
• 가중치 양자화(GPTQ, AWQ)와 KV 캐시 양자화는 별개 설정입니다.
• 흔한 실수 1: 가중치만 계산하고 KV 캐시를 빼먹는 것 — 롱 컨텍스트에서는 KV 가 지배적입니다.
• 흔한 실수 2: GiB 와 GB 를 섞는 것 — 이 실습은 2^30 기준입니다.

단계 7개

1. 공식 구현하기
2. 7B 4K 기준값으로 검증하기
3. 롱 컨텍스트로 확장하기
4. GQA 적용해 줄이기
5. KV 캐시 양자화 적용하기
6. GPU 에 들어가는 최대 동시 요청 수 구하기
7. 시나리오 표 생성하기