NVIDIA/kvpress

KVPress – 장문 LLM을 위한 KV 캐시 압축

개요

• 트랜스포머 기반 언어 모델의 키-값(KV) 캐시를 프리필링 단계(및 선택적으로 디코딩 중)에서 가지치기하거나 압축하는 다양한 프레스(압축 알고리즘)를 구현한 Python 라이브러리입니다. KV 캐시를 줄이면 메모리 사용량이 줄어들고 매우 긴 컨텍스트(예: 100만 토큰)에서 추론 속도가 빨라질 수 있습니다.

중요성

• 현대 LLM에서 KV 캐시는 처리된 토큰 수에 비례하여 증가합니다. 70B 모델을 100만 토큰 문서에서 실행하려면 300GB 이상의 GPU 메모리가 필요할 수 있습니다. KVPress를 사용하면 연구자와 개발자는 모델 코드를 다시 작성하지 않고도 최신 압축 방법을 많이 실험할 수 있습니다.

주요 기능

• 플러그 앤 플레이 파이프라인 – KVPressTextGenerationPipeline이 Hugging Face 파이프라인(kv-press-text-generation)으로 등록됩니다.
• 방대한 프레스 카탈로그 – 단순한 무작위 가지치기부터 정교한 어텐션 기반 및 행렬 분해 기법(예: KnormPress, SnapKVPress, KVzipPress, CURPress)까지 30개 이상의 즉시 사용 가능한 방법을 제공합니다.
• 프리필 및 디코딩 압축 – DecodingPress 및 관련 래퍼를 사용하여 토큰 생성 중에 주기적으로 압축할 수 있습니다.
• 구성 가능 – 래퍼 프레스(AdaKVPress, ComposedPress, ChunkPress 등)를 사용하여 전략을 쌓거나 결합할 수 있습니다.
• 훈련 불필요 – 나열된 모든 프레스는 추가 모델 미세 조정 없이 작동합니다.
• 양자화 캐시 지원 – Hugging Face의 QuantizedCache(예: optimum-quanto를 통한 4비트)와 함께 작동합니다.
• 멀티 GPU 지원 – accelerate 디바이스 맵과 호환됩니다.
• 벤치마크 및 리더보드 – 평가 CLI 및 노트북과 함께 RULER 4k 토큰 벤치마크의 공개 HF Space 리더보드를 제공합니다.

설치

pip install kvpress               # 기본 설치
# 또는 개발용으로 선택적 종속성(평가, flash-attention) 포함
uv sync --extra eval --extra flash-attn

기본 사용 예시

from transformers import pipeline
from kvpress import ExpectedAttentionPress

pipe = pipeline(
    "kv-press-text-generation",
    model="Qwen/Qwen3-8B",
    device_map="auto",
    dtype="auto",
)

context = "A very long text you want to compress once and for all"
question = "\nA question about the compressed context"
press = ExpectedAttentionPress(compression_ratio=0.5)

answer = pipe(context, question=question, press=press)["answer"]
print(answer)

프레스는 모델이 질문에 주의를 기울이기 전에 컨텍스트 토큰의 KV 캐시를 압축하여 다양한 압축 비율을 비교할 수 있게 합니다.

평가

• 각 프레스를 장문 컨텍스트 데이터셋에서 실행하고 정확도, 메모리 절감, 속도 향상을 보고하는 CLI를 제공합니다.
• 저장소에는 대화형 프로파일링용 노트북(speed_and_memory.ipynb, wikipedia_demo.ipynb)이 포함되어 있습니다.
• RULER 4k 토큰 벤치마크 결과는 README에 링크된 Hugging Face 리더보드에 표시됩니다.

문서 및 리소스

• Colab 노트북 데모, HF Space 데모, 블로그 게시물을 통해 라이브러리를 소개합니다.
• 각 프레스는 소스 파일과 해당 arXiv 논문에 대한 링크를 제공하여 더 깊이 학습할 수 있습니다.
• "새 프레스" 노트북은 사용자 정의 압축 방법을 추가하는 방법을 안내합니다.

라이선스

• Apache 2.0 (README의 배지 참조).

인용

@article{devoto2025expectedattention,
  title={Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution},
  author={Devoto, Alessio and Jeblick, Maximilian and J{\'e}gou, Simon},
  journal={arXiv preprint arXiv:2510.00636},
  year={2025},
  url={https://arxiv.org/abs/2510.00636}
}

획득처

• GitHub: https://github.com/NVIDIA/kvpress
• PyPI: kvpress
• Hugging Face Space: https://huggingface.co/spaces/nvidia/kvpress