KV cache와 추론 비용

KV 캐시가 긴 문맥 추론을 비싸게 만들어

제곱 비용이 주로 학습에서 드러난다면, 추론에서 더 직접적으로 마주치는 문제는 KV 캐시야. Autoregressive generation은 이미 처리한 모든 토큰의 key와 value 텐서를 저장해 뒀다가 새 토큰을 만들 때 다시 써. 캐시는 문맥 길이에 비례해 커지므로 문맥이 길어질수록 GPU 메모리의 큰 몫을 차지해.

숫자로 보면 문제가 더 선명해져. 매개변수 70B, 어텐션 head 64개, head 차원 128, 80개 층으로 이뤄진 fp16 서비스 모델은 문맥 토큰 하나마다 대략 2 × n_layers × n_heads × head_dim × 2 bytes를 써. 128K 토큰에서는 모델 가중치와 별도로 캐시만 40GB 이상 필요해. H100 한 장에 사용자 세션 하나를 담기도 빠듯하고, 여러 요청을 묶어 처리하는 효율도 무너지지.

GQA와 MQA로 KV 메모리를 줄여

Grouped Query Attention(GQA, Llama 2/3에서 사용)은 여러 query head가 K와 V head를 함께 쓰게 해. Llama 2 70B는 query head 64개와 KV head 8개를 사용해 품질 손실을 거의 내지 않으면서 KV 캐시를 8분의 1로 줄였어. Multi-Query Attention(MQA)은 KV head를 하나만 둬 메모리를 더 아끼지만, 어려운 작업에서는 측정 가능한 품질 손실이 생겨. 그래서 오늘날 최전선 모델 대부분은 GQA를 기본값으로 삼아.

그래도 문맥에 따라 선형으로 커지는 성질은 남아 있어. 상수 계수를 낮췄을 뿐 기울기는 그대로야. GQA를 써도 문맥 128K 모델의 서비스 비용은 8K보다 요청당 약 3–5배 비싸. 긴 문맥을 제공하는 업체가 별도 API 가격과 가산 요금을 매기는 데는 이런 구조적 비용이 반영되어 있어.

Code

흔한 모델 구성에서 KV 캐시 크기(GB)·python

def kv_cache_gb(seq_len, n_layers, n_kv_heads, head_dim, dtype_bytes=2):
    # Per token: 2 (K and V) * layers * heads * head_dim * dtype_bytes
    bytes_per_tok = 2 * n_layers * n_kv_heads * head_dim * dtype_bytes
    return seq_len * bytes_per_tok / (1024 ** 3)

# Llama 3.1 70B는 GQA: 8 KV heads, 80 layers, head_dim 128
for n in [8_000, 32_000, 128_000, 1_000_000]:
    print(f'{n:>9} tok -> {kv_cache_gb(n, 80, 8, 128):>6.2f} GB KV-cache')
# 8000   -> 2.50 GB
# 32000  -> 10.00 GB
# 128000 -> 40.00 GB
# 1000000 -> 312.50 GB

Exercise

배포 후보인 세 모델 계열의 토큰당 KV 캐시 크기를 계산해 봐. Llama 3.1 8B(GQA, KV head 8개, 32개 층, head 차원 128), Llama 3.1 70B(8, 80, 128), 가상의 MHA 70B(64, 80, 128)를 비교해. 문맥 길이 4K부터 1M까지 KV 캐시만 차지하는 GPU 메모리를 그래프로 그리고, H100 80GB 한 장에 더는 들어가지 않는 지점을 표시해. 서비스 팀이라면 반드시 확인해야 할 그림이야.

KV 캐시가 긴 문맥 추론을 비싸게 만들어

GQA와 MQA로 KV 메모리를 줄여

Code

External links

Exercise

Progress

댓글 0