KV-Cache 양자화를 위한 RoPE 인지적 비트 할당

기존의 저비트 KV-cache 양자화기(quantizers)는 각 캐시된 키(key)를 평면 벡터(flat vector)로 취급하는 경우가 많습니다. 그러나 RoPE (Rotary Positional Embedding) 하에서, 미래의 어텐션 로짓(attention logit)에 대한 키의 기여도는 2차원 주파수 블록(frequency blocks)에 대한 위치 의존적 합계로 분해됩니다. 이는 키-캐시 양자화를 블록 단위의 비트 할당(bit-allocation) 문제로 만듭니다. 즉, 에너지가 높은 RoPE 블록은 양자화 오차에 더 민감하므로 더 많은 비트를 할당받아야 합니다. 우리는 TurboQuant-MSE(TQ-MSE)를 기반으로 구축된 키-캐시 양자화를 위한 RoPE 인지적 비트 할당기인 Block-GTQ를 소개합니다. 각 레이어(layer)와 KV 헤드(head)에 대해, Block-GTQ는 각 RoPE 블록에 대한 레이블이 없는 에너지 점수(label-free energy score)를 계산하고, 한계 이득(marginal gain)에 따라 정수 비트 폭(integer bit widths)을 탐욕적(greedily)으로 할당합니다. 일치된 K/V 비트 예산 하에서, Block-GTQ는 10개의 모델 진단 패널에서 RoPE 쿼리-키 로짓(query-key logits)을 더 잘 보존하며, 2 및 3 b/dim K-only 양자화 시 레이어당 MAE(Mean Absolute Error)를 32-80%까지 절감하고, 균일한(uniform) TQ-MSE와의 367개 레이어 비교 전 항목(367/367)에서 승리했습니다. 이러한 충실도(fidelity) 향상은 더 강력한 다운스트림 긴 문맥 검색(long-context retrieval), 이해 및 추론으로 이어집니다. Llama-3.1-8B-Instruct의 K2V2 환경에서 Block-GTQ는 6개 작업의 NIAH(Needle In A Haystack) 평균을 70.6에서 97.4로, LongBench-EN 평균을 36.87에서 53.31로 높였습니다. DeepSeek-R1-Distill-Qwen-7B를 사용한 AIME 2024/2025에서 fp16 최근-키 버퍼(recent-key buffer) 없이 K3V2 설정의 Block-GTQ는 51.7/37.5의 점수를 기록하여 fp16의 54.2/37.9에 근접한 반면, 균일한 TQ-MSE는 0.0/0.0으로 무너졌습니다. 우리는 추가로 팩트-캐시(packed-cache) 서빙 경로를 구현했습니다. Qwen2.5-3B-Instruct를 사용하는 단일 H800 GPU에서, 팩트 K3V3는 fp16과 비교 가능한 품질로 3.24배의 KV-cache 압축을 달성하며, 128K 문맥에서 fp16 FlashAttention2보다 1.34배 빠르게 실행되고, 피크 메모리(peak memory)를 56.31 GB에서 19.85 GB로 줄이며, fp16이 OOM(Out of Memory)이 발생하는 256K 및 512K에서도 실행 가능합니다. 코드는 https://github.com/JIA-Lab-research/blockgtq 에서 확인할 수 있습니다.