WUSH-KV: 데이터 적응형 변환으로 저비트 KV 캐시 양자화 개선

배경

대규모 언어 모델의 긴 컨텍스트 추론에서 KV 캐시는 핵심적인 시스템 병목입니다. 새로운 토큰을 생성할 때마다 이전 토큰의 Key와 Value를 읽어야 하므로, 컨텍스트 길이와 배치 크기가 증가하면 캐시 메모리와 메모리 대역폭 부담도 함께 커집니다. KV 캐시를 낮은 비트로 저장하면 이 비용을 줄일 수 있지만, 지나친 양자화는 어텐션 계산의 정확도를 떨어뜨릴 수 있습니다.

WUSH-KV는 단순히 비트 수를 줄이는 대신, 양자화 전에 데이터를 더 적합한 형태로 바꾸는 접근을 취합니다. 핵심은 실제 캘리브레이션 데이터의 통계를 활용해 제한된 양자화 수준을 효율적으로 사용하는 것입니다.

주요 아이디어

• WUSH의 KV 캐시 적용: WUSH는 행렬 곱을 이루는 두 인자의 2차 통계에서 데이터 인식형 변환을 만들어 양자화 오류를 줄입니다. WUSH-KV는 이 원리를 KV 캐시에 맞게 확장했습니다.
• Key와 Value를 분리 처리: 캘리브레이션 데이터를 이용해 Key와 Value에 서로 다른 변환을 구성합니다. 어텐션에서 두 요소가 담당하는 역할의 차이를 반영하기 위한 설계입니다.
• 계산 그래프에 맞춘 배치: Value 변환은 모델 가중치에 접어 넣을 수 있고, Key 변환은 RoPE 적용 이후에 수행됩니다. 이를 통해 기존 어텐션 흐름과의 결합을 고려합니다.
• 클리핑 양자화와 결합: 변환은 특정 양자화 방식에만 한정되지 않습니다. 논문은 QuEST INT와 조합했을 때, 일정한 가정 아래 WUSH 변환이 최적에 가까운 특성을 보인다고 분석합니다.

평가 결과

QuEST INT를 사용한 비교에서 WUSH-KV는 레이어별 재구성 오류를 줄였고, 테스트된 변환 가운데 가장 낮은 엔드투엔드 perplexity를 기록했다고 보고됐습니다. 또한 연구진은 WUSH-KV를 SGLang에 통합하고, OSCAR 스타일의 백분위수 클리핑 아핀 양자화를 사용해 서빙 환경에 가까운 평가를 진행했습니다.

2비트 설정에서는 평가된 모든 모델과 다운스트림 작업에서 WUSH-KV가 OSCAR 변환과 비슷하거나 더 나은 성능을 보였습니다. 다만 제공된 소재에는 구체적인 모델명, 작업 목록, 절대 수치가 제시되지 않았습니다. 따라서 이 결과를 모든 조건에서의 절대적 우위로 해석하기보다는, 여러 평가 조건에서 일관된 경쟁력을 보였다는 의미로 보는 편이 적절합니다.

의미와 남은 과제

WUSH-KV의 차별점은 양자화를 어텐션 구조와 함께 설계한다는 데 있습니다. Key와 Value의 역할, RoPE의 적용 위치, 가중치에 변환을 흡수하는 방법을 함께 고려해 제한된 비트를 더 유용하게 배분하려고 합니다. 이는 긴 컨텍스트 추론의 효율을 높이려면 비트 수뿐 아니라 데이터 표현 방식도 중요하다는 점을 보여줍니다.

실제 배포에는 과제도 남아 있습니다. 캘리브레이션 데이터와 오프라인 변환 생성이 필요하고, 가중치 통합 및 런타임 연산을 서빙 프레임워크가 지원해야 합니다. 실제 이득은 하드웨어가 저비트 연산과 변환 오버헤드, 캐시 접근을 얼마나 효율적으로 처리하는지에 따라서도 달라질 수 있습니다.

출처: Hugging Face Daily Papers