STAR-KV: 적응형 랭크 제어를 위한 소프트 임계값 기반의 저차원 (Low-Rank) KV 캐시 압축

저차원 투영 (Low-rank projection)은 은닉 차원 (hidden-dimension)의 중복성을 활용하여 KV 캐시를 압축하는 유망한 접근 방식으로 부상했습니다. 그러나 기존 방법들은 고정되거나 휴리스틱한 (heuristic) 랭크 선택에 의존하며, 정확도 저하를 최소화하면서 공격적인 압축을 달성하는 데 어려움을 겪고 있습니다. 우리는 미세한 랭크 제어가 가능한 적응형 저차원 KV 캐시 압축 프레임워크인 STAR-KV를 제안합니다. STAR-KV는 1) 어텐션 헤드 (attention-head) 및 블록 (block) 수준 모두에서 최적의 랭크 선택을 가능하게 하는 미분 가능한 임계값 메커니즘 (differentiable thresholding mechanism), 2) 키 (key) 및 값 (value) 투영의 민감도에 따라 서로 다른 저차원 분해를 적용하는 하이브리드 분해 전략 (hybrid decomposition strategy), 3) 데이터 통계량을 활용하여 손실에 가까운 저비트 양자화 (near lossless low-bit quantization)를 수행하는 저차원 인지 혼합 정밀도 양자화 (low-rank-aware mixed precision quantization)를 포함합니다. 여러 LLM 및 벤치마크를 통해 평가한 결과, STAR-KV는 최대 75%의 KV 캐시 압축을 달성하였으며, 양자화와 결합 시 최대 20배의 전체 KV 캐시 감소를 보여주었습니다. 커스텀 Triton 기반 GPU 커널을 통해 구현된 STAR-KV는 어텐션 모듈에서 최대 6.9배의 속도 향상과 3.1배의 엔드투엔드 (end-to-end) 생성 처리량 (throughput)을 제공합니다. 우리의 코드는 다음에서 공개적으로 사용할 수 있습니다: https://github.com/PriyanshBhatnagar/STAR-KV.