KVmix: 레이어 중요도 기반 혼합 정밀도 KV 캐시 양자화

초록

KVmix은 LLM 추론 시 KV 캐시의 메모리 부담을 완화하기 위해, 각 레이어의 키·밸류 투영 가중치에 대한 그래디언트 L2 노름을 중요도 지표로 활용한다. 중요도가 높은 레이어에는 높은 비트폭을, 낮은 레이어에는 저비트폭을 할당하는 혼합 정밀도 양자화를 수행하고, 최근 피벗 토큰은 풀프리시전으로 유지하면서 오래된 토큰은 Aggressive하게 압축한다. 실험 결과 Llama·Mistral 모델에서 평균 2.2‑2.4비트 수준의 KV 저장을 달성하면서 4.9배 메모리 절감·5.3배 추론 속도 향상을 기록한다.

상세 분석

본 논문은 KV 캐시가 LLM의 자동 회귀 디코딩에서 차지하는 메모리 비중이 시퀀스 길이에 비례해 선형적으로 증가한다는 점을 출발점으로 삼는다. 기존 정적 양자화 기법은 모든 레이어에 동일한 비트폭을 적용해 메모리‑정확도‑처리량 사이의 트레이드오프를 피할 수 없으며, 동적 양자화 기법은 높은 연산 비용과 KV 중요도 고려 부족으로 실시간 적용이 제한적이었다. KVmix은 이러한 한계를 극복하기 위해 두 가지 핵심 아이디어를 제시한다.

첫째, 그래디언트 기반 레이어 중요도 분석이다. 키와 밸류 투영 행렬 (W_k^i, W_v^i)에 대해 손실 함수 (L)에 대한 그래디언트 (\nabla_{W_k^i}L, \nabla_{W_v^i}L)의 L2 노름을 계산한다. 이 값은 양자화 오차 (\Delta K, \Delta V)가 손실에 미치는 1차 근사 (\Delta L \approx \langle \partial L/\partial K, \Delta K\rangle + \langle \partial L/\partial V, \Delta V\rangle)에서 계수 역할을 하므로, 노름이 클수록 해당 레이어의 KV가 모델 출력에 미치는 영향이 크다고 판단한다. 논문은 다중 프롬프트에 대해 평균을 취해 안정적인 중요도 점수 (\bar{s}_k^i, \bar{s}_v^i)를 얻고, 이를 기반으로 레이어를 상위 20 %와 하위 80 %로 구분한다(비율은 사용자 요구에 따라 조정 가능).

둘째, 동적 피벗 컨텍스트 선택이다. 중요도 분석 결과를 활용해 최근 토큰(피벗 컨텍스트)과 오래된 토큰을 구분한다. 피벗 토큰은 모델이 현재 토큰을 예측할 때 가장 큰 기여를 하므로 풀프리시전(16‑bit)으로 유지하고, 오래된 토큰은 레이어별 중요도에 따라 2‑3 bit 수준으로 압축한다. 이렇게 하면 긴 시퀀스에서도 핵심 정보 손실을 최소화하면서 메모리 사용량을 크게 줄일 수 있다.

구현 측면에서는 비대칭 저비트 양자화와 CUDA 커스텀 커널을 설계해 양자화·복호화 연산의 오버헤드를 최소화한다. 특히 3‑bit 양자화에서 스케일·오프셋을 레이어별로 다르게 적용해 정밀도 손실을 보정한다. 프로파일링 단계는 오프라인으로 한 번 수행되며, 이후 추론 시 추가 비용이 발생하지 않는다.

실험에서는 Llama‑2‑7B, Llama‑2‑13B, Mistral‑7B 등 여러 모델에 KVmix을 적용했으며, KV 비트 평균이 Key 2.19 bit, Value 2.38 bit에 불과함에도 불구하고 GSM8K, TruthfulQA, WikiText 등 다양한 벤치마크에서 FP16 대비 0.1 % 이하의 정확도 저하만을 보였다. 메모리 압축률은 평균 4.9×, 추론 처리량은 5.3× 향상되었으며, 기존 QA‑Q, KVTuner와 비교해 동일 메모리 제한 하에서 1‑2 % 이상의 정확도 우위를 확보했다.

본 연구는 레イヤ별 중요도와 시퀀스 시간축을 동시에 고려한 혼합 정밀도 양자화라는 새로운 패러다임을 제시한다. 그래디언트 기반 중요도 측정은 가중치 자체가 정적인 특성을 활용해 빠르고 비용 효율적인 프로파일링을 가능하게 하며, 동적 피벗 전략은 긴 컨텍스트 상황에서도 품질 저하 없이 메모리를 절감한다. 향후 연구에서는 중요도 측정에 헤드‑단위 혹은 토큰‑단위 세분화를 도입하거나, 스파스 양자화와 결합해 더욱 극단적인 메모리 절감 효과를 탐색할 여지가 있다.