Zefan-Cai/KVCache-Factory

해결하는 문제

KVCache-Factory는 긴 문맥을 가진 대규모 언어 모델(LLM) 추론 시 키-값(KV) 캐시와 관련된 메모리 한계를 해결합니다. 문맥 길이가 길어질수록 KV 캐시는 선형적으로 증가하여 막대한 GPU 메모리를 소비하고 추론 속도를 저하시킵니다. 이 프로젝트는 모델 성능을 희생시키지 않고 메모리 사용량을 줄이기 위한 다양한 전략을 구현, 비교, 평가할 수 있는 통합 프레임워크를 제공합니다.

작동 방식

이 프로젝트는 여러 KV 캐시 최적화 기법을 하나의 평가 인터페이스 아래 통합하는 '실험실' 역할을 합니다. 다음과 같은 여러 유형의 방법을 지원합니다:

압축 및 제거: 중요도가 낮은 토큰을 제거 (예: StreamingLLM, H2O, SnapKV, PyramidKV).
검색: 현재 쿼리에 가장 관련 있는 토큰만 선택 (예: Quest, L2Norm).
병합: 여러 KV 쌍을 하나로 집계 (예: CAM).
양자화: 캐시된 값의 정밀도를 낮춤 (예: KIVI, KVQuant, GEAR).
오프로딩: KV 캐시를 CPU 메모리로 이동하고 GPU로 하나씩 스트리밍하여 GPU 메모리 오버플로우 없이 정밀도를 유지 (HeadInfer).

LongBench, RULER, "Needle-in-a-haystack" 등의 표준 벤치마크를 위한 러너를 포함하여 이러한 방법이 정확도와 지연에 미치는 영향을 측정할 수 있습니다.

대상 사용자

이 도구는 LLM 효율성에 관심 있는 AI 연구자 및 개발자, 특히 긴 문맥 창 성능 최적화 또는 다양한 KV 캐시 관리 전략을 벤치마크하려는 사람들을 위한 것입니다.

주요 특징

통합 인터페이스: PyramidKV, SnapKV, H2O, StreamingLLM 등 다양한 베이스라인을 지원합니다.
다양한 최적화 전략: 제거, 검색, 병합, 양자화, 손실 없는 오프로딩을 하나의 플랫폼에서 통합합니다.
포괄적인 벤치마킹: LongBench, RULER, Needle-in-a-haystack 평가를 내장 지원합니다.
하드웨어 유연성: 대규모 모델(예: Llama-3-70B)용 멀티 GPU 추론을 지원하며, FlashAttention v2 및 SDPA 경로를 제공합니다.
시각화 도구: 주의 패턴을 시각화하여 다양한 압축 방법이 모델의 집중도에 어떻게 영향을 미치는지 이해할 수 있는 유틸리티를 포함합니다.