[논문 리뷰] KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization

Coleman Hooper, Sehoon Kim외 5명|arXiv (Cornell University)|2024. 01. 31.

Natural Language Processing Techniques인용 수 9

한 줄 요약

KVQuant은 채널별, pre-RoPE, 비균일, 벡터 단위 밀집-희소 KV 캐시 양자화와 Q-Norm을 도입하여 초장문 맥락 추론(최대 10M)에서 최소한의 perplexity 손실과 측정 가능한 속도 향상을 가능하게 한다.

ABSTRACT

LLMs are seeing growing use for applications which require large context windows, and with these large context windows KV cache activations surface as the dominant contributor to memory consumption during inference. Quantization is a promising approach for compressing KV cache activations; however, existing solutions fail to represent activations accurately in sub-4-bit precision. Our work, KVQuant, facilitates low precision KV cache quantization by incorporating several novel methods: (i) Per-Channel Key Quantization, where we adjust the dimension along which we quantize the Key activations to better match the distribution; (ii) Pre-RoPE Key Quantization, where we quantize Key activations before the rotary positional embedding to mitigate its impact on quantization; (iii) Non-Uniform KV Cache Quantization, where we derive per-layer sensitivity-weighted non-uniform datatypes that better represent the distributions; and (iv) Per-Vector Dense-and-Sparse Quantization, where we isolate outliers separately for each vector to minimize skews in quantization ranges. By applying our method to the LLaMA, Llama-2, Llama-3, and Mistral models, we achieve < 0.1 perplexity degradation with 3-bit quantization on both Wikitext-2 and C4, outperforming existing approaches. Our method enables serving LLaMA-7B with a context length of up to 1 million on a single A100-80GB GPU and up to 10 million on an 8-GPU system. We develop custom CUDA kernels for KVQuant, showing that we can achieve up to ~1.7x speedups, compared to baseline fp16 matrix-vector multiplications, for the LLaMA-7B model.

연구 동기 및 목표

• 장문 컨텍스트 LLM 추론에서 KV 캐시 메모리 병목을 줄이려는 동기를 부여한다.
• 길어진 시퀀스 및 배치에서 KV 캐시 활성화가 메모리의 대부분을 차지함을 보인다.
• 최소한의 perplexity 손실로 초저정밀도를 달성하기 위한 양자화 기법의 집합을 제안한다.
• LLaMA, LLaMA-2, Mistral 모델에서 긴 컨텍스트 평가를 포함한 엔드-투-엔드 개선을 시연한다.
• 실용적 배치를 위한 커널 차원의 구현과 오픈 소스 도구를 제공한다.

제안 방법

• KV 캐시 분포를 분석하여 이상치와 RoPE 효과를 식별한다.
• Key에 대해 채널별 양자화, Value에 대해 토큰별 양자화를 적용한다.
• RoPE로 인한 분포 이동을 완화하기 위해 Key를 RoPE 이전에 양자화한다.
• 보정(calibration)을 통해 오프라인으로 도출된 민감도 가중 비균일 데이터타입 nuqX를 개발한다.
• 벡터당 이상치를 분리하기 위한 벡터별 밀집-희소 양자화를 도입한다.
• 양자화 후 중심값을 양자화 전 통계와 일치시키도록 Q-Norm을 제안한다.
• 추론 중 KV 양자화 및 역양자화를 가속하기 위해 커스텀 CUDA 커널을 구현한다.

실험 결과

연구 질문

• RQ1 대형 LLM에서 KV 캐시 활성화를 sub-4-bit 정밀도로 양자화하더라도 무시할 만한 perplexity 손실이 발생하는가?
• RQ2 채널별 및 pre-RoPE 양자화 전략이 기존 방법보다 Key 분포를 더 잘 보존하는가?
• RQ3 KV 캐시 정확도에 대해 비균일(nuqX) 양자화와 벡터별 이상치 처리의 효과는 어떠한가?
• RQ4 다중 GPU 구성에서 초장문 컨텍스트 추론(최대 10M 토큰)에 대해 어떤 메모리 절감과 속도 향상이 달성 가능한가?
• RQ5 KVQuant의 가중치 양자화 및 엔드-투-엔드 배포와의 호환성은 어떠한가?

주요 결과

방법	LLaMA-7B 혼란도	LLaMA-7B KV 캐시 (GB)	LLaMA-13B 혼란도	LLaMA-13B KV 캐시 (GB)	LLaMA-30B 혼란도	LLaMA-30B KV 캐시 (GB)	LLaMA-65B 혼란도	LLaMA-65B KV 캐시 (GB)
baseline	5.68	32.0	5.09	50.0	4.10	97.5	3.53	160.0
int4	5.98	8.0	5.32	12.5	4.34	24.4	3.73	40.0
nf4	5.87	8.0	5.23	12.5	4.25	24.4	3.63	40.0
int4-gs128	5.77	8.3	5.16	13.0	4.16	25.3	3.57	41.6
int4-gs64	5.73	8.6	5.14	13.5	4.14	26.3	3.56	43.1
nf4-gs128	5.77	8.5	5.17	13.3	4.17	25.9	3.58	42.5
nuq4	5.73	8.0	5.15	12.5	4.16	24.4	3.57	40.0
nuq4-1%	5.70	8.6	5.11	13.5	4.12	26.3	3.54	43.2
int3	10.87	6.0	8.69	9.4	6.82	18.3	6.37	30.0
nf3	7.33	6.0	6.21	9.4	5.46	18.3	4.44	30.0
int3-gs128	6.17	6.3	5.47	9.8	4.44	19.2	3.78	31.5
int3-gs64	5.93	6.6	5.29	10.3	4.26	20.1	3.66	33.0
nf3-gs128	6.26	6.5	5.52	10.2	4.54	19.8	3.83	32.5
nuq3	6.01	6.0	5.34	9.4	4.41	18.3	3.74	30.0
nuq3-1%	5.76	6.6	5.15	10.4	4.17	20.3	3.59	33.2
int2	11779	4.0	69965	6.3	1470	12.2	7272	20.0
nf2	3210	4.0	5786	6.3	2044	12.2	5367	20.0
int2-gs128	37.37	4.3	41.77	6.7	16.49	13.0	13.63	21.4
int2-gs64	11.09	4.6	9.84	7.1	6.60	13.9	5.54	22.8
nf2-gs128	351.23	4.5	141.19	7.0	60.97	13.7	31.69	22.5
nuq2	8.17	4.0	7.29	6.3	7.05	12.2	27.17	20.0
nuq2-1%	6.06	4.6	5.40	7.3	4.43	14.2	3.76	23.2

• 3-bit KV 캐시 양자화는 LLaMA-7B/13B/30B/65B에서 Wikitext-2 및 C4에 대해 0.1 이하의 perplexity 저하를 야기한다.
• KV 캐시 메모리 사용량이 약 4.8배 감소한다.
• 커스텀 CUDA 커널로 Key/Value 매트릭스-벡터 연산의 최대 1.4배 속도향상.
• 장문 컨텍스트 추론(단일 A100-80GB에서 최대 1M; 8 GPUs에서 10M)이 KVQuant로 가능해진다.
• Pre-RoPE Key 양자화 및 채널별 전략이 기존 방법을 크게 능가한다.
• 벡터별 밀집-희소 양자화 및 1% 이상치 제거를 통한 이상치 처리로 거의 FP16 수준의 성능을 유지한다.
• Q-Norm은 양자해제된 분포를 양자화 전 통계와 맞춰 2-bit 양자화 정확도를 향상시킨다.