KV Cache Quantization
마지막 수정:
KV cache는 weight도 아니고 일반 activation도 아니다.
token이 들어올 때 K/V를 한 번 계산한다
그 K/V를 cache에 저장한다
이후 decode step마다 계속 다시 읽는다
즉 생성되는 순간은 activation처럼 동적이지만, 저장된 뒤에는 serving state처럼 오래 남는다.
2K
4.0x users- FP16 KV
- 296
- INT4 KV
- 1184
8K
4.0x users- FP16 KV
- 74
- INT4 KV
- 296
32K
4.1x users- FP16 KV
- 18
- INT4 KV
- 74
128K
4.5x users- FP16 KV
- 4
- INT4 KV
- 18
KV cache가 커지는 이유
Autoregressive decoding에서는 매 step마다 과거 token을 다시 참고한다. 매번 과거 token의 key와 value를 다시 계산하면 너무 비싸기 때문에 cache에 저장한다.
layers
heads
sequence length
head dim
K and V
batch / concurrent requests
이 항목들이 곱해지면서 KV cache가 커진다. Context length가 길어질수록 cache는 선형으로 늘고, batch나 concurrent users가 늘어도 같이 커진다.
그래서 weight를 4-bit로 줄여도 long-context serving에서는 KV cache가 다시 GPU memory를 잡아먹는다.
KV cache quantization의 첫 번째 이득은 capacity다
KV cache를 INT8이나 INT4로 줄이면 같은 GPU memory에 더 많은 cache block을 담을 수 있다.
more cache capacity
-> longer context
-> larger batch
-> more concurrent users
-> more prefixes cached
하지만 이것이 곧바로 per-token latency 감소를 뜻하지는 않는다.
만약 attention kernel이 quantized KV를 다시 FP16으로 dequantize한 뒤 high-precision attention을 수행한다면, FLOPS 자체가 크게 줄지 않는다. 이 경우 이득은 latency보다 memory capacity와 throughput 여유에 가깝다.
Latency까지 줄이려면 quantized attention kernel이 같이 필요하다.
KV cache quantization only
-> smaller cache
-> more capacity
KV cache quantization + quantized attention kernel
-> smaller cache + lower-precision attention path
Keys와 values는 같은 방식으로 quantize하면 안 된다
Key와 value는 같이 저장되지만 attention 안에서 역할이 다르다.
Key:
query와 dot product되어 어떤 token을 볼지 결정한다
Value:
attention weight로 섞여서 어떤 정보를 전달할지 결정한다
이 차이 때문에 통계 구조도 다르다.
책의 TinyLlama 분석에서는 key가 channel-wise outlier를 보인다. 같은 hidden channel이 여러 token과 prompt에서 반복적으로 커진다.
keys
-> channel-wise structure
-> per-channel scale이 유리
Value는 token-wise variation이 더 중요하다. 어떤 token의 value vector 전체가 커질 수 있다.
values
-> token-wise structure
-> per-token scale이 유리
그래서 좋은 기본값은 asymmetric KV granularity다.
K: per-channel quantization
V: per-token quantization
Scale overhead는 보통 감당할 만하다
Per-token value scale은 sequence length만큼 늘어난다. 그래서 metadata overhead가 생긴다.
하지만 head dimension이 128인 예시에서 INT4 value data는 token당 64 bytes이고, FP16 scale은 token당 2 bytes 정도다.
64 bytes data
2 bytes scale
-> about 3% overhead
정확도를 지키기 위한 비용으로는 작은 편이다.
더 낮은 bit로 가면 vector quantization이 등장한다
INT4 정도에서는 scalar quantization과 key/value별 granularity가 실용적인 기본값이다.
하지만 3-bit 이하로 내려가면 균일한 scalar quantization이 더 어려워진다. 이때 TurboQuant 같은 방법은 KV vector를 online으로 압축한다.
책의 TurboQuant 흐름은 다음과 같다.
1. KV vector의 norm과 direction을 분리한다
2. random rotation으로 outlier energy를 퍼뜨린다
3. fixed codebook으로 각 coordinate를 quantize한다
이 방식은 calibration data 없이 들어오는 vector를 즉시 압축할 수 있다는 장점이 있다. 다만 runtime과 kernel 지원 측면에서 일반적인 INT4 KV cache quantization보다 더 특수한 주제다.
이 path에서는 먼저 INT4 KV cache와 key/value granularity를 기본으로 이해하면 된다.
적용 순서
실무적으로는 보통 weight 또는 weight+activation quantization을 먼저 본다.
그 다음 문제가 long context나 decode-heavy throughput이면 KV cache quantization을 추가한다.
1. weight-only or W8A8/FP8로 model path 최적화
2. long context / high concurrency에서 KV memory 확인
3. KV cache quantization 적용
4. quantized attention kernel 지원 여부 확인
5. quality와 throughput을 함께 측정
KV cache quantization은 weight quantization과 경쟁하지 않는다. 서로 다른 tensor를 줄이므로 같이 쓸 수 있다.
AWQ/GPTQ weights
+ INT4 or FP8 KV cache
-> model memory and cache memory both reduced
확인
- KV cache quantization이 latency보다 capacity/concurrency에 먼저 영향을 주는 이유는 무엇인가?
- Key는 per-channel, value는 per-token scale이 유리한 이유는 무엇인가?
- KV cache quantization만으로 부족하고 quantized attention kernel이 필요한 경우는 언제인가?