KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache

논문 원본: https://arxiv.org/abs/2402.02750

KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache

논문 한 줄 요약: KIVI는 Key와 Value의 통계적 구조 차이를 이용 -> Key는 per-channel, Value는 per-token으로 비대칭 2bit 양자화를 적용함으로써 정확도 손실 없이 KV 캐시 메모리와 추론 비용을 크게 줄이는 방법

 

발표 자료:

엣지컴퓨팅 7조.pdf

5.26MB

---

문제 인식

LLM 추론에서 KV Cache는

• 시퀀스 길이와 배치 크기에 비례해 급격히 커지고
• 매 토큰마다 메모리에서 다시 로드되어
• 메모리와 속도 모두의 병목이 된다.

따라서 KV Cache 자체를 줄이는 것이 가장 직접적인 해결책!

 

문장이 처음 들어오는 Prefill 단계와 토큰을 한 개씩 생성하는 Decoding 단계로 이루어져있다.

논문은 여러 모델과 레이어에서 KV Cache 분포를 분석했고, 다음 차이를 발견

• Key Cache
  • 특정 channel에 고정된 outlier 존재
  • 레이어가 달라도 동일한 채널에서 반복됨
  • per-token 양자화 시 큰 오차 발생
• Value Cache
  • outlier 거의 없음
  • token-wise로 균일한 분포
  • attention에서 여러 토큰이 섞여 사용됨

→ Key와 Value는 같은 방식으로 양자화하면 안 됨