NVIDIA가 KVPress 툴킷으로 KV Cache 압축 기법을 제공해 1M 토큰 처리 시 메모리 사용량 327.6GB에서 유의미한 수준으로 감소

Context

자동회귀 모델의 장문 컨텍스트 처리 시 KV Cache 메모리 사용량이 컨텍스트 길이에 선형으로 증가하는 문제가 발생했다. Llama 3-70B를 bfloat16 정밀도로 1M 토큰 컨텍스트에서 실행하려면 KV Cache만 327.6GB 메모리가 필요해 실질적 배포가 불가능했다.

Technical Solution

• KV Cache 메모리 문제 해결을 위해 NVIDIA가 개발한 KVPress 툴킷 도입: 상태-최신-미술 KV Cache 압축 기법들을 Python 패키지로 제공
• 다양한 압축 알고리즘(Press) 적용: KnormPress(낮은 키 값 노름의 KV 쌍 제거), SnapKVPress(최신 쿼리에 대해 낮은 어텐션 가중치의 KV 쌍 제거)
• AdaKVPress와 ExpectedAttentionPress의 조합으로 RULER 데이터셋에서 최고 성능 달성
• Transformers 라이브러리의 KV Cache Quantization과 함께 통합 가능한 모듈식 설계
• 128K 컨텍스트 길이에서 50% 압축 비율 적용 시 메모리 45GB에서 37GB로 감소

Impact

• 128K 컨텍스트 길이에서 50% 압축 비율 적용 시 피크 메모리 사용량 45GB에서 37GB로 18% 감소
• A100 GPU에서 디코딩 속도 11 tokens/sec에서 17 tokens/sec로 55% 향상
• RULER, InfiniteBench, Loogle 벤치마크 데이터셋에 대한 9가지 압축 기법의 성능 비교 및 측정 가능

Key Takeaway

KV Cache 압축은 선형적 메모리 확장 문제의 실질적 해결책이지만, 높은 압축 비율은 모델 정확도 손실이라는 트레이드오프를 수반한다. 장문 컨텍스트 LLM 배포 시 압축 비율과 정확도 간 균형을 찾기 위해 벤치마크 기반의 체계적 평가가 필수적이다.