피드로 돌아가기
Dev.toAI/ML
원문 읽기
KV Cache 최적화를 통한 1.3TB OOM 위험 제거 및 메모리 효율 8배 개선
KV Cache Is Eating Your VRAM — Here's How to Estimate It Before You Run Out
AI 요약
Context
LLM 추론 시 정적 모델 가중치보다 동적 KV Cache가 VRAM 점유의 핵심 병목으로 작용함. 특히 Context Length와 Batch Size의 선형적 증가로 인해 고성능 GPU 환경에서도 예측 불가능한 OOM 발생 가능성이 높음.
Technical Solution
- GQA(Grouped Query Attention) 도입을 통한 KV Head 수 축소로 메모리 점유율 1/8 수준으로 절감
- KV Cache 전용 Quantization(FP16 → FP8/INT4) 적용을 통한 파라미터당 바이트 수 감소 및 정밀도 손실 최소화
- PagedAttention 기반의 비연속적 메모리 할당 체계 구축으로 메모리 파편화(Fragmentation)를 1~4% 수준으로 억제
- Sliding Window Attention 적용을 통해 특정 윈도우 외부의 토큰 캐시를 제거함으로써 메모리 증가 상한선 설정
- max_model_len 설정을 통한 실제 워크로드 기반의 Context Length 제한으로 불필요한 VRAM 낭비 차단
- Speculative Decoding 도입 시 Draft 모델과 Target 모델 각각의 KV Cache 합산량을 고려한 메모리 설계
Impact
- Llama 3.1 70B 모델 기준 MHA 대비 GQA 적용 시 토큰당 캐시 메모리 2.6MB에서 0.33MB로 감소
- vLLM의 PagedAttention 활용 시 기존 대비 30~50%의 유효 메모리 확보 가능
- KV Cache Quantization(INT4) 적용 시 7B 모델(8K, batch=16) 기준 메모리 사용량 18GB에서 4.5GB로 절감
실천 포인트
1. 배포 전 [2 * layers * hidden_dim * context_len * bytes_per_param] 공식으로 KV Cache 최대치 계산
2. vLLM 등 PagedAttention 지원 프레임워크 우선 채택으로 메모리 파편화 제거
3. FP8/INT4 KV Cache Quantization 적용 검토
4. GQA 아키텍처 모델 선정 및 워크로드에 최적화된 max_model_len 설정