Kitty: Accurate and Efficient 2-bit KV Cache Quantization with Dynamic Channel-wise Precision Boost
Created by
Haebom
Category
Empty
저자
Haojun Xia, Xiaoxia Wu, Jisen Li, Robert Wu, Junxiong Wang, Jue Wang, Chenxi Li, Aman Singhal, Alay Dilipbhai Shah, Alpay Ariyak, Donglin Zhuang, Zhongzhu Zhou, Ben Athiwaratkun, Zhen Zheng, Shuaiwen Leon Song
개요
Kitty는 LLM 추론의 주요 메모리 병목 현상인 KV 캐시 문제를 해결하기 위해 개발된 알고리즘-시스템 공동 설계 기법입니다. 4비트 KV 양자화는 정확도를 유지하지만, 2비트 양자화는 특히 긴 컨텍스트 추론에서 정확도를 저하시킵니다. Kitty는 Dynamic Channel-wise Precision Boost (민감도에 따라 Key-cache 채널을 랭킹하고, 소수의 채널만 더 높은 정밀도로 유지) 알고리즘을 사용하여 정확도 손실을 거의 없이 2비트에 가까운 메모리 사용량을 달성합니다. 페이지 레이아웃 통합, 균일한 역양자화, 분산된 읽기 및 하드 코딩된 마스크 문제 해결을 위해, Kitty는 혼합 정밀도 Key 페이지를 통일된 2비트 정밀도의 두 개의 텐서로 분해합니다. 이를 기반으로 페이지 중심 KV 레이아웃, Triton 호환 페이지 역양자화 커널 및 통합을 유지하고 발산을 방지하는 가벼운 런타임 파이프라인을 제공합니다. Qwen3 및 LLaMA3 모델을 사용하여 7가지 작업에서 테스트한 결과, Kitty는 KV 메모리를 거의 8배 줄이면서 정확도 손실을 무시할 수 있으며, 동일한 메모리 예산에서 최대 8배 더 큰 배치 및 2.1배에서 4.1배 더 높은 처리량을 가능하게 합니다.
시사점, 한계점
•
시사점:
◦
Dynamic Channel-wise Precision Boost 알고리즘을 통해 2비트 메모리 사용량에 가깝게 접근하면서도 정확도 저하를 최소화했습니다.
◦
혼합 정밀도 KV 캐싱을 위한 페이지 중심 레이아웃, Triton 호환 커널, 가벼운 런타임 파이프라인을 제공하여 성능을 최적화했습니다.