NVIDIA/kvpress

LLM KV cache compression made easy

What it solves

긴 컨텍스트 대형 언어 모델(LLM)을 배포하는 비용이 높은 이유는 키‑밸류(KV) 캐시가 토큰 수에 비례해 선형적으로 증가하기 때문입니다. 예를 들어, 70B 모델이 float16으로 100만 토큰을 처리하면 최대 330GB의 메모리가 필요할 수 있습니다. KVPress는 KV 캐시를 압축함으로써 이러한 메모리 오버헤드를 감소시키고 디코딩 속도를 향상시킵니다.

How it works

KVPress는 다양한 중요도 점수에 기반해 KV 쌍을 가지치기하거나 병합하는 "presses" 라는 압축 방법들을 제공하는 라이브러리입니다. 이러한 presses는 커스텀 KVPressTextGenerationPipeline을 통해 Hugging Face transformers 파이프라인에 통합됩니다.

  • Prefilling Compression: 입력 컨텍스트의 초기 처리 단계에서 캐시를 압축합니다.
  • Decoding Compression: 토큰 생성 중에 주기적으로 캐시를 압축하여 목표 크기를 유지합니다.
  • Scorer-based Pruning: KnormPress 혹은 SnapKVPress와 같은 많은 방법이 KV 쌍에 점수를 부여하고, 중요도가 가장 낮은 것을 제거합니다.
  • Quantization Support: QuantizedCache와 연동되어 4비트와 같은 저비트 정밀도로 메모리를 추가로 절감합니다.

Who it’s for

긴 컨텍스트 LLM 효율성을 연구·개발하는 연구자와 개발자, 특히 KV 캐시 압축 기술을 구현·벤치마크·배포하여 GPU 메모리 사용량을 줄이고자 하는 사람들을 위한 프로젝트입니다.

Highlights

  • Extensive Library: StreamingLLM, SnapKV, PyramidKV, KVzip 등 훈련 없이 사용할 수 있는 다양한 압축 방법을 구현했습니다.
  • Flexible Integration: 컨텍스트 매니저나 커스텀 transformers 파이프라인으로 쉽게 적용할 수 있습니다.
  • Hybrid Strategies: PrefillDecodingPress를 통해 프리필 단계와 디코딩 단계에 서로 다른 압축 방법을 조합할 수 있습니다.
  • Evaluation Tools: CLI와 노트북이 포함되어 있어 정확도, 피크 메모리 사용량, 전체 시간 절감 효과를 측정할 수 있습니다.