RoPE 인식 기반 비트 할당을 통한 KV-캐시 양자화 (RoPE-Aware Bit Allocation for KV-Cache Quantization)

긴 컨텍스트 추론에서 KV 캐시가 메모리·대역폭 병목을 유발하므로, 캐시를 저비트로 안전하게 압축하면 더 긴 컨텍스트를 단일 GPU에서 실행 가능하게 만든다. RoPE 구조를 반영한 블록 단위 비트 할당으로 key-side logit 보존을 우선하면, 동일한 평균 비트 예산에서 기존 균등 할당보다 attention 품질과 다운스트림 태스크 성능이 크게 개선된다.

Transformer의 RoPE(rotaional positional encoding)는 d_h 차원을 2차원 주파수 블록들로 나누고, 상대 위치 Δ에 대해 각 블록에 2×2 회전 R(Δθ_i)를 적용한다. 이에 따라 query-key 로그릿은 블록별 내적들의 합으로 정확히 분해되며, 각 블록의 기여는 쿼리·키 블록 노름과 블록 주파수에 의존한다. 따라서 키 오류가 로그릿에 미치는 영향은 블록별로 독립적으로 평가 가능하다. 기존 KV-cache 양자화는 키를 평탄한 벡터로 보고 전체 복원 오차를 줄이는 데 초점을 맞춘다. 그러나 RoPE 환경에서는 일부 주파수 블록이 로그릿 신호의 대부분을 차지해(높은 energy) 동일 비트 할당이 고에너지 블록을 과소보호하고 저에너지 블록을 과보호하는 결과를 낳는다. 이 차이는 attention 분포와 downstream 성능으로 이어진다. Block-GTQ는 각 RoPE 블록의 라벨-프리 에너지 점수 s_i(=½(E[||q^{(i)}||^2]+E[||k^{(i)}||^2]))를 계산해, TurboQuant-MSE의 로컬 rate law(복원 MSE ∝ 4^{-b})와 결합한 목적 J(b)=Σ s_i 4^{-b_i}을 최소화하는 정수 비트 할당을 수행한다. 마진 이득 Δ_i = (3/4) s_i 4^{-b_i}를 기준으로 한 비트씩 그리디로 할당하면 정수 제약 하에서 최적 해를 얻는다. 이 접근은 동일 평균 비트 예산 하에서 고에너지 블록에 비트를 집중시켜 query-key 로그릿 MAE를 크게 낮춘다. 결과적으로 softmax 기반 attention 분포 보존이 향상되며, NIAH/LongBench-EN/AIME 같은 장문 검색·생성·추론 작업에서 fp16에 가까운 성능을 더 낮은 메모리·대역폭으로 달성한다.

전체 접근: 각 레이어·KV 헤드에 대해 RoPE 블록별 에너지 점수를 계산하고, 정수 비트 예산 B(헤드 단위)에 대해 블록별 비트 b_i를 할당한다. 목적함수는 J_{9l,h}(b)=Σ s_{9l,h,i} 4^{-b_i}이며, b_min≤b_i≤b_max와 Σ b_i=B 제약을 둔다. 값(V)은 로그릿 계산에 관여하지 않으므로 uniform TQ-MSE로 처리한다. 블록 에너지 산출: s_{9l,h,i}=½(E_{t,g∈G(h)}[||q_{9l,g,t}^{(i)}||2^2]+E_t[||k{9l,h,t}^{(i)}||_2^2]) 형태로, 라벨이 없는 짧은 캘리브레이션 프리픽스에서 Q/K 두 번째 모멘트를 평균해 계산한다. G(h)는 KV 헤드 h를 읽는 쿼리 헤드 집합이다. 그리디 할당 알고리즘: 모든 블록을 b_min으로 초기화(B_extra=B-L b_min). 각 블록의 현재 마진 Δ_i=(3/4) s_i 4^{-b_i}를 우선순위 큐에 넣고, B_extra가 남는 동안 Δ_i가 최대인 블록에 비트 1을 추가한다. 비트가 증가하면 해당 블록의 다음 Δ는 4배 작아지므로 마진 기반 균형이 자동으로 유지된다. 알고리즘은 이 목적에 대해 정수 최적임을 보였다. 물리적 인코딩과 그룹화: 동일 비트 폭을 할당받은 블록들을 그룹 {G_b}로 묶어 해당 그룹을 하나의 TQ-MSE 인코더로 처리한다. 이렇게 하면 RoPE-블록 단위 결정은 유지되면서 작은 블록마다 별도 인코더를 만들지 않아도 된다. Packed-cache 서빙 경로: 양자화된 코드 스트림(니블 혹은 바이트)과 per-group codebook(공유 fp16), 노름 및 메타데이터를 HBM에 저장한다. fused attention 커널이 타일 단위로 코드+노름을 로드해 L1에 작은 fp16 코드북을 유지하며, 디코드 후 fp16 입력·fp32 누적 텐서코어 연산으로 QK^T와 P V를 계산한다. 전체 fp16 KV를 HBM에 전개하지 않아도 되므로 per-step HBM 트래픽은 코드+노름(예: K3V3에서 약 157 B/토큰/헤드)으로 줄어든다.

메인 벤치마크: Ten-model diagnostic panel에서 K-only( V는 fp16)로 평가한 결과, Block-GTQ는 동일 평균 비트 예산에서 per-layer RoPE-logit MAE를 32%–80% 감소시켰고(모델별 평균치), 3-bit 및 2-bit 예산 모두에서 모든 레이어 비교 367/367에서 uniform TQ-MSE를 이겼다(Table 1). 장문 검색·이해: Llama-3.1-8B-Instruct의 NIAH 다중태스크(4K–128K 범위, 여러 needle 깊이)에서 K2V2 예산 기준으로 TQ-MSE의 평균 70.6% 대비 Block-GTQ는 97.4%로 상승했다(Table 4(a)). Qwen2.5-7B-Instruct에서는 TQ-MSE가 0.0으로 붕괴한 반면 Block-GTQ는 60% 수준을 유지했다. 생성·추론(Reasoning): AIME 2024/2025 연산(DeepSeek-R1-Distill-Qwen-7B, K3V2)에서 no-buffer 설정(최근 토큰 버퍼 없음) 하에 Block-GTQ는 fp16(54.2/37.9)과 근접한 51.7/37.5을 기록했다. 동일 설정에서 uniform TQ-MSE는 0.0/0.0으로 성능이 붕괴했다(Table 6). 배포·효율성: Qwen2.5-3B-Instruct, H800, K3V3에서 packed-cache 경로로 실행 시 KV-cache 압축률 3.24×, T=128K에서 fp16 FlashAttention2 대비 decode 속도 1.34×, 피크 GPU 메모리 56.31GB→19.85GB로 감소했다. 또한 fp16가 OOM 되는 256K/512K 컨텍스트에서도 실행이 가능했다.

아키텍처 구조 개요: 각 Transformer 레이어의 KV head는 d_h 차원을 L=d_h/2개의 2차원 RoPE 블록으로 분할한다. Block-GTQ는 각 레이어·KV 헤드별로 블록 에너지 s_{9l,h,i}를 캘리브레이션 샘플에서 계산한 뒤 정수 비트 배분 b_i를 산출해, 동일 b를 가진 블록들을 그룹으로 묶어 TQ-MSE로 인코딩한다. 수학적 기반과 최적화 목표: 목표함수는 J(b)=Σ_{i=1}^L s_i 4^{-b_i}다. TurboQuant-MSE의 rate law E||x-9hat{x}||^2 ≤ ||x||^2 C_{TQ} 4^{-b}가 근거로 쓰인다. 한 비트를 추가할 때의 마진 감소는 Δ_i(b_i) = s_i 4^{-b_i} - s_i 4^{-(b_i+1)} = (3/4) s_i 4^{-b_i}로 닫힌 형태로 계산되며, 이 마진을 기준으로 그리디 할당을 수행하면 정수 제약 하에서 전역 최적을 얻는다는 정리를 보였다(증명은 Appendix A.4). 블록 에너지 산출 세부: s_{9l,h,i}는 캘리브레이션 프리픽스 상의 Q/K 두 번째 모멘트 평균으로 계산된다: s_{9l,h,i}=½(E_{t,g∈G(h)}[||q_{9l,g,t}^{(i)}||^2]+E_t[||k_{9l,h,t}^{(i)}||^2]). 여기서 G(h)는 해당 KV 헤드를 읽는 쿼리 헤드 집합이며, 캘리브레이션은 라벨이 없는 짧은 토큰 시퀀스(예: WikiText-2 일부)로 수행된다. Packed-cache 구현 디테일: K 스트림은 mixed-rate 스케줄로 낮은 비트 그룹은 nibble 컨테이너, 높은 비트 그룹은 바이트로 저장한다. fused attention 커널은 타일 단위로 코드를 로드해(니블 추출 포함) 작은 fp16 코드북을 L1에 상주시킨 뒤, 복원된 fp16 입력을 fp32로 누적해 QK^T와 PV를 계산한다. 긴 컨텍스트 처리는 키 축 분할로 병렬화하고 섹션별 log-sum-exp 병합을 통해 정확도를 유지한다. Prefill은 레이어별로 전체 프롬프트를 한 번에 quantize해 O(T) 호출을 O(1) 호출로 줄인다. 하드웨어/소프트웨어 구현: 단일 H800에서 Triton 커널로 K/V를 패킹·전송하고, fused attention은 tensor-core 기반 fp16 입력·fp32 누적으로 구현했다. Packed-path는 단일-스텝 디코딩에서 코드 디코딩 오버헤드를 지불하므로 짧은 컨텍스트에서는 fp16에 늦지만, KV 대역폭이 지배적인 긴 컨텍스트에서는 속도·메모리 이득이 역전된다.

논문에서 언급한 한계만 기술함: (1) Block-GTQ는 K에 대해서만 블록별 비트 할당을 수행하고 V는 uniform TQ-MSE로 남겨 두었으므로, V 민감한 백본에서는 최근-토큰 버퍼가 없을 때 성능 저하가 관찰됐다. (2) Packed-cache 경로의 공개 구현은 단일-GPU 단일-스레드 기반이며 multi-GPU·배치 서빙 확장은 향후 과제로 남아 있다. (3) Block-GTQ는 mixed-rate 메타데이터로 인해 동일 압축률의 균일 TQ-MSE보다 약간 더 많은 메타데이터 오버헤드가 있다.