TurboQuant: KV Cache Vector Quantization
마지막 수정:
TurboQuant는 GPTQ나 AWQ처럼 model weight를 줄이는 방법이 아니다.
TurboQuant의 target은 decode 중 계속 쌓이고 계속 읽히는 KV cache다.
GPTQ / AWQ:
offline으로 weight를 INT4에 가깝게 저장
TurboQuant:
inference 중 생성되는 KV vector를 online으로 압축
이 차이가 중요하다. Weight는 학습 이후 고정되어 있으므로 offline calibration이나 reconstruction을 할 수 있다. KV cache는 prompt와 생성 과정에 따라 매번 새로 만들어진다.
왜 scalar quantization만으로 부족한가?
INT8이나 INT4 KV cache quantization에서는 보통 scale을 잡고 각 값을 scalar로 quantize한다.
value -> scale로 나눔 -> integer grid에 맞춤
하지만 더 낮은 bit로 내려가면 coordinate 하나하나를 균일한 grid에 맞추는 방식이 급격히 어려워진다. 특히 KV vector 안에 outlier coordinate가 있으면 scale이 커지고, 나머지 coordinate의 표현력이 줄어든다.
TurboQuant는 여기서 관점을 바꾼다.
각 coordinate를 따로 양자화할 것인가?
보다 정확한 질문:
vector 전체의 모양을 더 압축하기 쉬운 좌표계로 바꿀 수 있는가?
TurboQuant의 세 단계
TurboQuant 흐름은 세 단계로 볼 수 있다.
1. Split
KV vector를 length와 direction으로 나눈다.
2. Rotate
큰 coordinate 하나에 몰린 energy를 여러 축으로 퍼뜨린다.
3. Codebook
회전된 작은 vector 조각을 가장 가까운 codeword index로 저장한다.
첫째, KV vector를 norm과 direction으로 나눈다. Norm은 vector의 크기이고, direction은 정보가 향하는 방향이다.
둘째, random rotation을 적용한다. Rotation은 vector의 길이를 보존하면서 coordinate 축을 바꾼다. 이때 한 coordinate에 몰려 있던 outlier energy가 여러 coordinate로 퍼지면, 낮은 bit codebook으로 표현하기 쉬워진다.
셋째, 회전된 vector 조각을 fixed codebook의 index로 저장한다. 즉 값을 그대로 저장하는 대신, 가장 가까운 대표 vector의 번호를 저장한다.
raw KV vector
-> norm + direction
-> rotated direction
-> codebook index sequence
Calibration이 아니라 online compression에 가깝다
TurboQuant는 calibration data로 static range를 정하는 방식과 성격이 다르다.
KV cache는 생성 중에 계속 들어온다. 따라서 TurboQuant는 들어오는 vector를 그 자리에서 회전하고 codebook에 맞춰 압축한다.
이 점 때문에 TurboQuant는 long-context나 high-concurrency에서 의미가 크다.
context length 증가
-> KV cache 증가
-> GPU memory 압박
-> 더 aggressive한 KV compression 필요
다만 실제 이득은 runtime과 kernel 지원에 달려 있다. KV cache를 작게 저장해도 attention kernel이 그것을 효율적으로 읽고 복원하지 못하면 latency 이득은 제한될 수 있다.
확인
- TurboQuant가 GPTQ/AWQ와 달리 weight quantization이 아닌 이유는 무엇인가?
- Random rotation이 outlier coordinate 문제를 완화하는 직관은 무엇인가?
- KV cache compression에서 runtime kernel 지원을 같이 봐야 하는 이유는 무엇인가?