WUSH-KV: 데이터 적응 변환을 이용한 KV 캐시 양자화

KV 캐시는 컨텍스트 길이와 배치 크기가 커짐에 따라 메모리 및 대역폭 비용이 증가하며, 이는 효율적인 장문맥 추론을 제한하는 요인이 됩니다. 이 병목 현상을 해결하기 위해, 우리는 저비트 KV-캐시 양자화를 위한 WUSH-KV를 소개합니다. WUSH-KV는 행렬 곱의 두 인자의 2차 통계량으로부터 데이터 인식 변환(data-aware transform)을 구성하여 양자화 오차를 줄이는 WUSH를 적용합니다. WUSH-KV는 보정 데이터를 사용하여 별도의 키 및 값 변환을 구성하며, 값 변환은 모델 가중치에 통합되고 키 변환은 RoPE 이후에 적용됩니다. 이 변환들은 클리핑된 양자화기(clipped quantizers)와 결합될 수 있습니다. 이러한 양자화기 중 하나인 QuEST INT를 사용하여, 우리는 온건한 가정 하에서 WUSH 변환이 근사 최적임을 보여줍니다. 이 양자화기를 사용함으로써, WUSH-KV는 레이어별 재구성 오차를 줄이고 테스트된 다른 변환들 중에서 가장 낮은 종단 간(end-to-end) 퍼플렉서티를 달성합니다. 종단 간 평가를 위해, 우리는 OSCAR 스타일의 분위수 클리핑 아핀 양자화(percentile-clipped affine quantization)를 사용하여 WUSH-KV를 SGLang에 통합했습니다. 2비트에서, WUSH-KV는 평가된 모든 모델과 다운스트림 작업 전반에 걸쳐 OSCAR 변환과 비교할 만하거나 능가하는 성능을 보였습니다.