llm-inference context-length memory-optimization kv-cache-quantization cuda-kernels per-channel-quantization pre-rope-quantization non-uniform-quantization
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.
한국어 요약
한 줄 요약
KVQuant은 3비트 정밀도로 0.1 perplexity 이하의 정확도 손실을 유지하면서 1000만 토큰 길이의 LLM 추론을 가능하게 하는 KV 캐시 양자화 기법이다.
핵심 기여도
- Per-Channel Key Quantization: 채널 단위로 Key 활성화를 양자화하여 분포에 더 잘 맞춤.
- Pre-RoPE Key Quantization: RoPE 적용 전 Key를 양자화하여 RoPE의 영향을 완화.
- Non-Uniform KV Cache Quantization (NUQ): 레이어별 민감도를 고려한 비균일 양자화 타입 도입.
- Per-Vector Dense-and-Sparse Quantization: 벡터 단위로 이상치를 분리하여 양자화 범위 왜곡 최소화.
핵심 아이디어
기존의 KV 캐시 양자화는 분포의 이상치와 RoPE의 영향으로 인해 4비트 이하에서는 정확도가 급격히 저하되었다. KVQuant은 Key 활성화의 채널 단위 분포를 분석하고, RoPE 전에 Key를 양자화함으로써 RoPE가 채널 간 분산을 일으키는 문제를 해결한다. 또한, NUQ는 민감도 기반 비균일 양자화를 통해 적은 수의 양자화 점을 최적 배치하고, Per-Vector 방식은 벡터 단위로 이상치를 분리하여 정확도를 유지한다. 특히, Q-Norm 레이어는 양자화 후 분포를 fp16 값과 일치시키는 역할을 하며, 2비트 양자화에서도 정확도를 개선한다.
기술적 접근법
- **Per-Channel Key Quantization**: Key 활성화의 채널 단위 분포를 고려한 양자화.
- **Pre-RoPE Key Quantization**: RoPE 전 Key 활성화를 양자화하여 RoPE의 혼합 효과를 방지.
- **Non-Uniform Quantization (NUQ)**: 레이어별 민감도를 고려한 비균일 양자화 타입을 오프라인에서 계산.
- **Per-Vector Dense-and-Sparse Quantization**: 벡터 단위로 이상치를 분리하여 1% 이상치 제거로 0.1 perplexity 이하의 정확도 유지.
- **Q-Norm Layer**: 양자화 후 분포를 fp16 값과 일치시키는 정규화 레이어.
- **CUDA 커널 최적화**: 활성화 양자화 과정을 효율적으로 수행하는 커스텀 CUDA 구현.
주요 결과
- **LLaMA-7B 모델**에서 3비트 양자화로 **Wikitext-2와 C4 데이터셋에서 0.1 perplexity 이하의 정확도 손실**을 기록.
- **1개 A100-80GB GPU**에서 **100만 토큰 길이**의 추론 가능.
- **8-GPU 시스템**에서 **1000만 토큰 길이**의 추론 가능.
- **4비트 정밀도에서 Key와 Value의 행렬-벡터 곱 연산에서 fp16 대비 약 1.4×의 속도 향상**.
의의 및 한계
KVQuant은 기존의 4비트 이하 양자화가 정확도를 유지하지 못하는 문제를 해결하여, 초장문 추론을 가능하게 한다. 특히, NUQ와 Per-Vector 방식은 이상치 처리와 분포 일치를 동시에 달성하며, Q-Norm은 2비트에서도 정확도를 개선한다. 그러나, 2비트 이하에서는 추가적인 정규화 레이어가 필요하며, 모든 모델에 동일한 성능이 보장되지 않을 수 있다. 또한, 이상치 제거 비율(1%)은 데이터셋에 따라 최적화가 필요하다.
실용적 활용
KVQuant은 대규모 문서 요약, 코드 분석, 장문 질의 응답 등 초장문 처리가 필요한 산업 및 연구 분야에 적용 가능하다. 특히, 고비용의 멀티GPU 환경 없이도 100만 토큰 길이의 추론을 가능하게 하므로, 클라우드 인프라 비용 절감에 기여할 수 있다.