왜 KV 캐시가 병목인가
트랜스포머 추론에서 이미 생성한 토큰의 Key/Value 벡터는 매 스텝 재사용되므로 캐시에 보관한다. 문제는 이 캐시 크기가 컨텍스트 길이에 선형으로 늘어난다는 점이다. 캐시 용량은 다음과 같이 계산된다.
bytes = 2 (K,V) × layers × seq_len × kv_heads × head_dim × dtype_bytes
# 예: 70B 급, 80 layers, kv_heads 8, head_dim 128, fp16, 32K 토큰
= 2 × 80 × 32768 × 8 × 128 × 2
≈ 10.7 GB (배치 1 기준)배치가 커지거나 컨텍스트가 128K로 늘면 캐시만으로 수십 GB를 소모해 가중치보다 캐시가 더 큰 상황이 벌어진다. 결국 동시 처리 가능한 요청 수(throughput)가 캐시 메모리에 묶인다.
양자화가 해결하는 지점
가중치 양자화와 달리 KV 캐시 양자화는 런타임에 생성되는 K/V 텐서를 fp16 대신 int8·fp8·int4로 저장한다. int8이면 캐시가 절반, int4면 1/4로 줄어 같은 GPU에서 컨텍스트를 2~4배 늘리거나 배치를 키울 수 있다. 연산 정밀도가 아니라 저장 정밀도만 낮추는 것이라 어텐션 계산 직전에 역양자화(dequant)한다.
정밀도별 트레이드오프
| 포맷 | 캐시 크기 | 품질 영향 | 권장 상황 |
|---|---|---|---|
| fp16 | 1.0x | 기준 | 여유 있을 때 |
| fp8 (e4m3) | 0.5x | 거의 없음 | 기본 권장 |
| int8 | 0.5x | 미미 | fp8 미지원 GPU |
| int4 | 0.25x | 과제 의존 | 초장문·요약 |
경험적으로 K는 이상치(outlier)가 많아 민감하고 V는 관대하다. 그래서 K는 int8, V는 int4처럼 비대칭 구성이 안전하며, per-token 또는 per-channel 스케일을 두면 int4에서도 열화를 크게 줄인다.
vLLM에서 적용하기
vLLM은 엔진 옵션 하나로 KV 캐시 양자화를 켠다. fp8은 별도 캘리브레이션 없이 바로 쓸 수 있다.
from vllm import LLM, SamplingParams
llm = LLM(
model="meta-llama/Llama-3.1-8B-Instruct",
kv_cache_dtype="fp8", # fp16 대비 캐시 50% 절감
max_model_len=131072, # 확보한 여유로 컨텍스트 확장
gpu_memory_utilization=0.90,
)
out = llm.generate(
["다음 문서를 요약하라: ..."],
SamplingParams(temperature=0.2, max_tokens=512),
)
print(out[0].outputs[0].text)서버 모드에서도 동일하게 플래그로 넘긴다.
vllm serve meta-llama/Llama-3.1-8B-Instruct \
--kv-cache-dtype fp8 \
--max-model-len 131072 \
--gpu-memory-utilization 0.9직접 구현할 때의 핵심
커스텀 추론 루프라면 스케일 계산과 clamp가 관건이다. per-token 대칭 int8 예시는 다음과 같다.
import torch
def quant_int8(x): # x: [tokens, heads, head_dim]
scale = x.abs().amax(dim=-1, keepdim=True) / 127.0
scale = scale.clamp(min=1e-8)
q = (x / scale).round().clamp(-127, 127).to(torch.int8)
return q, scale
def dequant_int8(q, scale):
return q.to(torch.float16) * scale
k_q, k_s = quant_int8(k_new) # 캐시에는 int8 + scale만 저장
k_full = dequant_int8(k_q, k_s) # 어텐션 직전 복원스케일을 텐서당 하나로 두면(per-tensor) 저장은 더 줄지만 이상치에 취약해진다. head_dim 축으로 스케일을 유지하는 편이 실패가 적다.
실무에서 주의할 점
- 품질 검증은 태스크로 한다. perplexity만 보지 말고 실제 사용 태스크(요약 정확도, 코드 정합성, 검색 재현율)로 A/B를 돌린다. 장문 검색(needle-in-haystack)에서 int4는 중간 위치 회수율이 떨어질 수 있다.
- 스케일 저장 오버헤드를 계산에 넣어라. per-token 스케일은 토큰마다 fp16 값이 추가되므로 head_dim이 작으면 절감폭이 기대보다 준다.
- dequant 지연. 양자화는 메모리를 아끼지만 역양자화 연산이 추가된다. 대개 메모리 대역폭 이득이 커서 순이득이지만, 짧은 컨텍스트·소규모 배치에서는 오히려 손해일 수 있다.
- GPU 지원 확인. fp8 e4m3는 Hopper(H100) 이상에서 하드웨어 가속된다. Ampere(A100)에서는 int8 경로가 현실적이다.
정리
KV 캐시 양자화는 가중치를 건드리지 않고 장문 처리의 메모리 상한을 직접 낮추는 수단이다. 우선 fp8로 시작해 절반을 확보하고, 더 필요하면 V부터 int4로 내리며 태스크 지표로 손실을 확인하는 순서가 안전하다. 캐시가 병목이었다면 같은 하드웨어에서 컨텍스트와 배치를 동시에 키울 수 있다.