PN-Mix가 RoPE 채널 쌍 중요도로 KV 캐시를 2.26비트까지 압축한다
PN-Mix: Pair-Norm Guided Mixed-Precision Key Cache Quantization for Efficient Long-Output LLM Inference
무엇인가
긴 출력을 생성하는 추론, 특히 코드 생성이나 복잡한 문제 풀이에서는 KV 캐시가 시스템 병목이 된다. 캐시는 시퀀스 길이와 배치 크기에 선형으로 커지면서 메모리 압박을 키우고, 서빙 가능한 배치 크기와 최대 생성 길이를 동시에 깎는다. 논문은 Qwen3-8B FP16 기준으로 단일 A100-40GB에서 배치 32일 때 생성 길이 4K를 넘기면, 또는 생성 길이 8K일 때 배치가 20을 넘기면 메모리가 부족해진다고 제시한다. KV 캐시 양자화는 토큰 위치와 캐시 구조를 그대로 둔 채 정밀도만 낮추는 해법이지만 손실 압축이어서, 2비트 같은 공격적 저비트 양자화는 긴 사고 사슬에서 추론 정확도를 크게 떨어뜨린다. 기존 혼합정밀도 방법들은 중요한 키 채널에 높은 비트를 주지만, 키 통계만으로 중요도를 추정하고 개별 채널 단위로 다뤄 query-key 상호작용과 RoPE가 만드는 채널 쌍 구조를 놓친다는 것이 저자들의 문제 제기다.
어떻게 동작하나
핵심 관찰은 attention logit 수준의 섭동 분석에서 나온다. RoPE는 인접한 두 채널을 한 쌍으로 묶어 회전시키므로, logit은 채널 쌍별로 ‖q_i‖·‖k_i‖·cos(φ)의 합으로 분해된다. 양자화로 키 쌍이 재구성되면 키 쌍 노름의 상대 섭동 ρ와 각도 항의 변화 δ가 생기는데, 이때 logit 오차의 기여도는 query-key pair-norm 곱이 스케일한다. 즉 어떤 채널 쌍의 양자화 오차가 logit을 흔드는지는 키만의 크기가 아니라 query와의 곱으로 결정된다. 저자들은 재구성 후 값이라 사전에 알 수 없는 ρ, δ 대신, 양자화 전 상태에서 바로 얻을 수 있는 이 곱을 민감도 대리 지표로 쓴다. 인과적 위치 쌍 전체에 대해 곱을 평균하는 대신, query 쌍 노름 기대값과 key 쌍 노름 기대값을 각각 구해 곱하는 분리 근사를 택했고, 이 근사가 인과 결합 점수의 중요 쌍 순위를 대체로 보존한다는 것을 Spearman 순위 상관과 top-k% 선택 중복으로 확인했다. 실제로 채널별 K̃ 범위, 채널별 K̃ 크기, K̃ 쌍 노름, 제안한 Q̃K̃ 쌍 노름 네 가지 기준을 같은 약 2.25비트 예산에서 비교했을 때, Q̃K̃ 쌍 노름이 두 모델 모두에서 attention 분포 KL 발산이 가장 낮았다.
무엇과 다른가
PN-Mix의 절차는 다음과 같다. prefill 단계에서 각 레이어·KV 헤드·채널 쌍마다 중요도 σ를 계산하는데, GQA에서는 하나의 KV 헤드가 여러 query 헤드와 공유되므로 해당 KV 헤드를 공유하는 query 헤드들의 query 쌍 노름을 평균해 곱한다. KV 헤드별로 σ 상위 r% 채널 쌍을 선별하고, 배치 크기 B에 대해 B×L×H^kv×(d/2) 비트짜리 이진 마스크로 저장한다. 이 선택은 prefill에서 한 번만 하고 디코딩 내내 재사용하므로, 오프라인 캘리브레이션 없이 입력 적응적이면서도 디코딩 중 중요도 재계산이나 그룹별 인덱스 저장·로딩이 필요 없다. 온라인 양자화는 키와 값 모두 처음 N_s=32개 sink 토큰과 크기 N=128의 잔여 버퍼를 FP16으로 유지하고, 버퍼가 차면 키를 per-channel 비대칭 그룹으로 양자화한다. 이때 선별된 채널 쌍은 b_H=3비트, 나머지는 b_L=1비트를 쓰며, 1비트 그룹에는 재구성 레벨을 양 끝점에서 bin 중점 쪽으로 옮기는 보정을 적용해 극단적 min-max 1비트 양자화의 편향을 줄인다. 값은 §4.1에서 확인한 낮은 민감도에 따라 per-token 비대칭 그룹으로 b_V 비트로 양자화한다.
어떻게 쓰나
실험은 lm-evaluation-harness로 수행했다. 일반 모델(Llama-3.1-8B-Instruct, Qwen3-8B/32B, Qwen2.5-14B-Instruct)에는 GSM8K 8-shot, MATH-Algebra 4-shot, HumanEval, GPQA-Diamond 5-shot을 4096토큰 생성 제한으로, 추론 모델(DeepSeek-R1-Distill-Llama-8B, Qwen3-8B, DeepSeek-R1-Distill-Qwen-7B)에는 AIME 2024/2025를 32768토큰 제한과 temperature 0.6, top-p 0.95, top-k 20 설정으로 평가했다. 비교 대상은 균일 정밀도 계열 KIVI, PolarQuant와 혼합정밀도 계열 ChanMix, Kitty다. PN-Mix는 두 운용점을 보고하는데 PN-Mix(1)은 r=50%, G=128로 평균 2.26비트, PN-Mix(2)는 r=40%, G=64로 평균 2.29비트다. 결과적으로 균일 2비트 KIVI는 네 모델 평균에서 FP16 대비 5.0~9.7포인트 하락한 반면, PN-Mix(1)은 평균 키 캐시 2.26비트로 네 모델 모두 FP16과 1포인트 이내를 유지했다. ChanMix와 Kitty보다 적은 비트를 쓰면서도 Llama·Qwen 8B~32B 전반에서 평균 점수가 동등하거나 높았고, 과제별로는 GSM8K에서 Kitty와 차이가 작은 편이며 HumanEval과 GPQA-Diamond에서 동등 이상인 경향을 보였다. 두 운용점 중에서는 더 적은 비트를 쓰는 PN-Mix(1)이 네 모델 모두 평균 점수가 높았다.
전제와 한계
AIME에서는 6개 모델-벤치마크 조합 중 5개에서 비-FP16 방법 가운데 가장 높은 평균 정확도를 기록했고, 예외는 AIME25의 R1-Qwen-7B였다. 값 캐시를 2비트로 양자화한 경우에도 같은 양상이 나타났다. 절제 실험에서는 sink 토큰 보존과 1비트 보정을 각각 제거하면 정확도가 떨어졌고 보정 제거의 하락 폭이 더 컸으며, Q̃K̃ 쌍 노름을 K̃만 쓰는 기준으로 바꾸면 평균 정확도가 특히 AIME24에서 낮아졌다. 디코딩 중 통계를 누적해 그룹마다 중요 쌍을 다시 고르는 온라인 재선택 변형은 성능이 더 낮았을 뿐 아니라 그룹별 인덱스 저장·로딩 부담까지 안았다. 비트율 파레토 실험에서 3.0, 2.5, 2.25, 2.0, 1.5비트를 훑었을 때 Kitty와 PN-Mix 모두 2.5~3.0비트에서는 FP16에 근접했지만, 2.25비트 아래에서는 PN-Mix가 더 완만하게 성능이 떨어졌다.
효율 실험은 단일 A100-SXM4-40GB에서 Qwen3-8B로 측정했다. 40GB 예산에서 PN-Mix는 배치 32로 생성 길이 32K를, 생성 길이 8K로 배치 110을 지원해 FP16 대비 각각 8배 긴 길이와 5.5배 큰 배치를 가능하게 했고, 최대 GPU당 처리량은 FP16의 2.37배에 도달했다. 프롬프트 16K에서 Kitty 대비 TTFT 차이는 최대 12ms로 prefill 예측 오버헤드가 무시할 수준이었고, 긴 CoT 궤적(B=1, P=141, L_gen=32K)에서 KV 캐시 길이에 따른 스텝당 지연 증가를 Kitty보다 8.8% 줄여 B=1 세 설정에서 디코딩 시간이 4.0~6.5% 낮아졌다. B=32, P=141, L_gen=32K에서는 PN-Mix만 37.12GB로 메모리에 들어갔다.
실무 관점에서 이 논문이 주는 시사점은 명확하다. KV 캐시를 2비트 부근까지 밀어야 하는 서빙 환경이라면, 채널별 키 크기나 범위 같은 키 단독 통계로 중요도를 정하는 대신 RoPE 채널 쌍 단위에서 query와 키의 노름 곱을 보는 편이 같은 비트 예산에서 attention 분포를 더 잘 보존한다. 또한 중요도 선별을 prefill 한 번으로 끝내고 이진 마스크만 들고 다니는 구조라, 디코딩 스텝마다 중요 채널을 재계산하거나 그룹별 인덱스를 관리하는 비용을 피할 수 있다. 다만 이득은 RoPE 기반 디코더 전용 구조와 GQA를 전제로 하며, 실제 적용 시에는 자신의 모델·프롬프트 분포에서 prefill 선별이 디코딩 구간까지 유효한지, 그리고 1비트 보정과 sink 토큰 보존이 자기 워크로드에서도 필요한지 확인해야 한다.
저자들이 밝힌 한계는 두 가지다. 첫째, PN-Mix의 중요도 점수가 RoPE 채널 쌍 위에 정의되므로 RoPE 기반 디코더 전용 LLM에 가장 직접적으로 적용되며, RoPE식 채널 쌍 구조가 없는 모델로 같은 정식화를 확장하려면 추가적인 적응이 필요하다. 둘째, 효율 결과는 현재의 Triton 구현으로 단일 A100-SXM4-40GB에서 측정한 것이며, 실제 지연 시간은 하드웨어, 커널 최적화, 배칭 전략, 서빙 인프라에 의존하므로 PN-Mix의 정확한 속도 향상 폭은 배포 환경마다 달라질 수 있다.