NVIDIA/kvpress

해결하는 문제

긴 컨텍스트 대규모 언어 모델(LLM)을 배포하는 것은 비용이 많이 듭니다. 키-밸류(KV) 캐시는 토큰 수에 비례하여 선형으로 증가하기 때문입니다. 예를 들어, Llama 3.1-70B를 float16으로 100만 토큰 처리할 경우 최대 330GB의 메모리가 필요할 수 있습니다. KVPress는 KV 캐시를 압축함으로써 이러한 메모리 요구를 줄이고 디코딩 속도를 높입니다.

작동 방식

KVPress는 프리필링 단계 또는 디코딩 단계 중 주기적으로 KV 쌍을 제거하거나 병합하는 다양한 '프레스'(압축 방법)를 구현합니다. Hugging Face transformers 라이브러리와 통합되어 토큰화 및 채팅 템플릿을 처리하는 커스텀 KVPressTextGenerationPipeline을 제공합니다.

압축은 다음과 같은 전략을 통해 이루어집니다:

• 스코어 기반 제거: 중요도 스코어(예: 무작위, 역노름, 또는 어텐션 가중치)를 사용하여 가장 중요도가 낮은 KV 쌍을 제거합니다.
• 구조적 압축: 특정 캐시 패턴을 유지합니다. 예를 들어, 레이어 간에 피라미드 형태의 크기를 유지하거나, 초기 및 최근 토큰만 유지(StreamingLLM).
• 차원 감소: 채널 어텐션 스코어 기반으로 키의 차원을 압축합니다.
• 고급 기법: CUR 분해, 컨텍스트 재구성(KVzip), 또는 제거된 토큰을 생존한 이웃 토큰에 병합하는 기법을 사용합니다.

대상 사용자

긴 컨텍스트 LLM의 효율성에 관심 있는 연구자 및 개발자로, KV 캐시 압축 방법을 구현, 벤치마킹, 배포하기 위한 표준화된 프레임워크가 필요한 분들입니다.

주요 특징

• 풍부한 라이브러리: SnapKV, StreamingLLM, KVzip 등 학습 없이 사용 가능한 다양한 압축 방법을 지원합니다.
• 유연한 파이프라인: transformers 파이프라인을 커스터마이징하여 간편한 통합과 평가를 가능하게 합니다.
• 이중 단계 압축: 프리필링 단계와 디코딩 단계(‘DecodingPress’를 통해) 모두에서 압축을 지원합니다.
• 양자화 지원: QuantizedCache와 호환되어 추가적인 메모리 절감이 가능합니다.
• 벤치마킹 도구: CLI와 노트북을 제공하여 정확도, 최대 메모리 사용량, 총 시간 절감을 측정할 수 있습니다.