[논문 리뷰] KIVI : Plug-and-play 2bit KV Cache Quantization with Streaming Asymmetric Quantization

Zirui Liu, Jiayi Yuan외 6명|arXiv (Cornell University)|2023. 01. 01.

Quantum-Dot Cellular Automata인용 수 6

한 줄 요약

KIVI는 대규모 언어 모델 추론에서 키-값(KV) 캐시에 대한 플러그 앤 플레이, 튜닝 불필요한 2비트 비대칭 양자화 방법을 제안한다. 키에는 채널별 양자화를, 값에는 토큰별 양자화를 적용한다. 이로 인해 최대 메모리 사용량이 2.6배 감소하고 배치 크기가 최대 4배로 증가하며, Llama, Falcon, Mistral 모델에서 정확도 손실 최소화로 2.35×–3.47× 높은 처리량을 달성한다.

ABSTRACT

Efficiently serving large language models (LLMs) requires batching of many requests to reduce the cost per request. Yet, with larger batch sizes and longer context lengths, the key-value (KV) cache, which stores attention keys and values to avoid re-computations, significantly increases memory demands and becomes the new bottleneck in speed and memory usage. Additionally, the loading of the KV cache causes the computational core to be idle, which limits the inference speed. A straightforward and effective solution to reduce KV cache size is quantization, which decreases the total bytes taken by KV cache. However, there is a lack of in-depth studies that explore the element distribution of KV cache to understand the hardness and limitation of KV cache quantization. To fill the gap, we conducted a comprehensive study on the element distribution in KV cache of popular LLMs. Our findings indicate that the key cache should be quantized per-channel, i.e., group elements along the channel dimension and quantize them together. In contrast, the value cache should be quantized per-token. From this analysis, we developed a tuning-free 2bit KV cache quantization algorithm named KIVI. With hardware-friendly implementation, KIVI can enable Llama, Falcon, and Mistral models to maintain almost the same quality while using $\mathbf{2.6 imes}$ less peak memory (including model weight). This reduction in memory usage enables up to $\mathbf{4 imes}$ larger batch size, bringing $\mathbf{2.35 imes \sim 3.47 imes}$ throughput on real LLM inference workload. The source code is available at https://github.com/jy-yuan/KIVI.

연구 동기 및 목표

• 높은 배치 크기와 긴 컨텍스트 길이에서 발생하는 KV 캐시로 인한 증가하는 메모리 및 속도 병목 문제를 해결하기 위해.
• 인기 있는 대규모 언어 모델들에서 KV 캐시의 요소 분포를 분석하여 양자화 과정의 과제와 한계를 이해하기 위해.
• 모델 정확도를 유지하면서 메모리 사용량을 극적으로 줄이는 하드웨어 우수한 튜닝 불필요한 2비트 양자화 방법을 개발하기 위해.
• 효율적인 KV 캐시 압축을 통해 실제 대규모 언어 모델 추론 워크로드에서 더 큰 배치 크기와 더 높은 처리량을 가능하게 하기 위해.

제안 방법

• 모델 재학습이나 미세조정이 필요 없는 플러그 앤 플레이 2비트 비대칭 양자화 방법인 KIVI를 제안한다.
• 키 캐시에 채널별 양자화를 적용하여 채널 차원에 따라 요소를 그룹화하고, 각 채널 내에서 양자화 오차를 제한한다.
• 값 캐시에 토큰별 양자화를 적용하여 자동회귀 생성의 스트리밍 특성과 일치시키고 오차를 각 토큰 단위로 격리한다.
• KV 캐시를 그룹화된 부분과 잔여 부분으로 분할한다: 그룹화된 세그먼트에는 그룹 기반 양자화를 적용하고, 잔여 부분은 정밀도 유지한다.
• 어텐션 계산 중에 그룹화된 부분과 잔여 부분을 결합하기 위해 타일 기반 행렬 곱셈을 사용하여 정확도를 유지한다.
• 프리필 및 디코딩 단계에서의 양자화 오버헤드를 최소화하기 위해 하드웨어 우수한 구현을 채택한다.

실험 결과

연구 질문

• RQ1 LLM에서 키 캐시와 값 캐시의 요소 분포는 어떻게 다른가? 이러한 차이가 양자화 전략에 어떤 영향을 미치는가?
• RQ2 왜 키 캐시에는 채널별 양자화가 더 효과적인가? 값 캐시에는 토큰별 양자화가 더 좋은가?
• RQ3 스트리밍 및 동적 성격을 지닌 KV 캐시를 고려할 때, 튜닝 없이도 2비트 양자화가 높은 모델 정확도를 유지할 수 있는가?
• RQ4 다양한 양자화 전략이 실제 LLM 추론 워크로드에서 메모리 사용량, 배치 크기, 처리량에 어떤 영향을 미치는가?
• RQ5 KIVI의 그룹화 및 잔여 부분 분할 메커니즘은 대규모 스트리밍 KV 캐시의 효율적 양자화를 가능하게 하면서 정확도를 어떻게 유지하는가?

주요 결과

• 키 캐시의 채널별 양자화는 채널 간 오차 전파를 줄이며, 고값 이상치 채널 존재로 인해 필수적이다.
• 값 캐시의 토큰별 양자화가 필요하다. 왜냐하면 값 캐시가 어텐션 계산에서 믹서 역할을 하기 때문이며, 오차는 각 토큰 단위로 격리되어야 한다.
• KIVI는 Llama-2-7B에서 최대 메모리 사용량을 2.6배 감소시키며, 다양한 벤치마크에서 거의 동일한 모델 품질을 유지한다.
• 이 방법은 최대 4배의 더 큰 배치 크기와 실제 LLM 추론 워크로드에서 2.35×에서 3.47×까지 높은 처리량을 달성한다.
• 절단 실험 결과, 그룹 크기와 잔여 길이가 성능에 측정 가능한 영향을 미치지만 관리 가능하며, 최적의 설정은 경험적으로 도출된다.
• KIVI는 Llama, Falcon, Mistral 등 다양한 모델에서 뛰어난 성능을 유지하며, 2비트 정밀도에서도 정확도 손실이 최소한이다.