MLA 캐시 두 경로의 오차 성질을 분석해 INT4 공동 압축을 가능케 한 QuantMLA
QuantMLA: Function-Aligned Dual-Path Quantization for Low-Bit MLA KV Caching
무엇인가
MLA(Multi-Head Latent Attention)는 키와 값을 유도하는 공유 저랭크 콘텐츠 latent와 분리된 RoPE 키, 두 개의 캐시 경로로 표현력 높은 멀티헤드 어텐션을 압축한다. 그런데 캐시 메모리는 여전히 컨텍스트 길이와 배치 크기에 선형으로 늘어난다. 기존 MLA 양자화 연구는 콘텐츠 캐시를 FP8로 줄이면서 RoPE 키 캐시는 BF16 같은 고정밀로 남겨 두는 것이 일반적이었다(FlashMLA, SnapMLA 등). RoPE 경로를 낮은 비트로 양자화했을 때 생기는 기능적 영향은 제대로 규명되지 않았고, 두 캐시를 동시에 저비트로 압축하는 시도는 사실상 비어 있었다.
어떻게 동작하나
이论文의 출발점은 "캐시 재구성 오차가 같으면 출력 왜곡도 비슷하다"는 통념을 깨는 실측이다. 저자들은 INT4 양자화가 만드는 오차 방향으로 한 경로씩 교란하고 캐시 NMSE를 0.01로 맞춘 뒤 출력 왜곡을 측정했다. DeepSeek-V2-Lite, DeepSeek-V3-Base, DeepSeek-R1, Kimi-K2-Instruct, LongCat-Flash-Lite, GLM-4.7-Flash 등 4개 패밀리 6개 모델(16B~1T 파라미터)에서 RoPE 경로의 모델 평균 응답은 콘텐츠 경로의 1.941배에서 6.864배에 달했다. 원인은 오차가 흘러가는 경로의 차이다. 콘텐츠 오차는 어텐션 매칭과 값 집계, 그리고 두 항의 상호작용까지 세 갈래로 출력에 도달하지만, RoPE 오차는 위치 로짓을 통한 매칭 한 갈래로만 들어간다. 그런데도 RoPE 쪽 왜곡이 더 큰 이유를 저자들은 국소 증폭 모델 G_b = T_b·A_b + O_b로 설명한다. T_b는 평균 연산자 민감도, A_b는 양자화 오차 에너지가 민감한 좌표에 얼마나 배분됐는지, O_b는 좌표 간 방향성 결합을 나타낸다. 즉 같은 에너지의 오차라도 민감한 방향에 정렬되면 훨씬 큰 왜곡을 만든다.
무엇과 다른가
QuantMLA의 방법은 이 비대칭을 경로별 변환 공간과 목적함수로 옮긴 것이다. 두 경로는 등가 재매개변수화에 서로 다른 제약을 건다. 콘텐츠 latent는 매칭과 집계에 모두 쓰이므로, 저자들은 직교 회전 R_C와 양수 채널 스케일 S_C로 캐시를 재매개변수화하고 이를 소비하는 투영 행렬을 W̃ = S_C⁻¹R_CᵀΓW로 보정한다. 무가중 RMSNorm이 회전 등변성을 가지므로 R_C는 오프라인에서 가중치에 완전히 융합되고, 온라인 변환 비용이 사라진다. RoPE 경로는 위치 회전과 교환해야 하므로 임의의 조밀 회전은 쓸 수 없다. 대신 RoPE 주파수 쌍마다 블록 대각 SO(2) 회전 R_P와 블록 대각 스케일 S_P를 두고, 쿼리와 키에 서로 역방향으로 적용하는 상호 변환 Q_P' = Q_P R_P S_P⁻¹, K_P' = K_P R_P S_P를 쓴다. 이 변환은 완전정밀 위치 점수를 그대로 보존하면서 RoPE 회전과 교환하므로 역시 오프라인 융합이 가능하다. 학습 목적함수도 경로별 오차 경로에 맞춘다. 콘텐츠 경로는 BF16 교사 모델 대비 출력 투영 후 어텐션 출력의 정규화 재구성 오차를 최소화해 매칭·집계·상호작용 항을 한꺼번에 잡는다. RoPE 경로는 위치 점수 Z_P = Q_P K_Pᵀ의 재구성 오차를 최소화하는데, 저자들은 위치 점수 교란이 출력에 미치는 영향을 τ²/4·(Σ_h‖V̄_h‖²)·Σ_h‖Ω⊙ΔZ_P,h‖²_F로 유계화하는 명제(Proposition 2)를 제시하고, 이 상한이 위치 QK 재구성을 직접 제어한다는 것을 보인다. 보정은 사전학습 가중치를 고정한 채 WikiText-2 128개 시퀀스로 300스텝 최적화하고, 분리된 32개 시퀀스에서 최적 체크포인트를 고른다.
어떻게 쓰나
실행 계층에서는 FlashMLA 디코드 경로 위에 네이티브 저비트 MLA 백엔드를 구현했다. 언패킹과 디양자화를 어텐션 계산 안에 직접 융합해 헤드별 키·값을 펼치거나 전체 히스토리 BF16 캐시를 실체화하지 않는다. 패킹된 데이터와 양자화 메타데이터를 재구성된 BF16 타일과 분리해 버퍼링하여 프리페치가 어텐션 계산과 겹치게 하고, RoPE QK 계산을 더 일찍 스케줄할 수 있게 했다. 콘텐츠와 RoPE 키를 모두 INT4로 쓸 때는 융합 라이터가 모든 토큰의 패킹 INT4 복사본과 함께 처음 4개 sink 토큰 및 최근 128개 토큰의 양자화 전 BF16 값을 유지하고, 토큰마다 정확히 하나의 표현만 어텐션에 쓴다. 최근 윈도우를 벗어난 비-sink 토큰은 추가 양자화 없이 기존 패킹 복사본으로 전환되며, 모든 유효 토큰이 하나의 전역 softmax에 참여해 전체 히스토리 어텐션을 보존한다.
전제와 한계
실험은 4개 MLA 패밀리를 두 벤치마크 묶음으로 평가한다. DeepSeek-V2-Lite와 Moonlight-16B-A3B는 상식 5종, MMLU, GSM8K, FDA·SWDE·SQuAD 검색 과제로, LongCat-Flash-Lite와 GLM-4.7-Flash는 GPQA-Diamond, MMLU, GSM8K, MATH500, AIME25, HumanEval, LiveCodeBench로 구성된 추론·코드 묶음으로 평가한다. 비교 대상은 RTN과, MLA 캐시 구조에 맞게 변형한 SmoothQuant†, QuaRot†이며 모두 토큰별 비대칭 아핀 양자화기에 그룹 크기 64를 쓴다. C4R4(콘텐츠 4비트, RoPE 키 4비트)에서 QuantMLA는 DeepSeek-V2-Lite와 Moonlight-16B-A3B에서 BF16 대비 각각 0.26%p, 0.13%p 이내를 유지했다. GSM8K에서 DeepSeek-V2-Lite는 35.25%로 RTN 16.98%, SmoothQuant† 24.56%, QuaRot† 30.71%를 크게 앞섰고, Moonlight-16B-A3B는 74.53%로 BF16과 동일했다. 콘텐츠 캐시를 INT2까지 내린 C2R4에서도 두 모델의 스위트 평균이 64.42%, 72.70%로 같은 정밀도의 가장 강한 베이스라인을 각각 7.20점, 10.84점 앞선다. 추론·코드 묶음에서도 C4R4 기준 LongCat-Flash-Lite 62.01%, GLM-4.7-Flash 53.75%로 최강 동일정밀도 베이스라인을 4.82점, 2.70점 상회했고, LiveCodeBench(LongCat)와 GSM8K(GLM)는 BF16 대비 각각 0.29점, 0.15점 차이에 그쳤다. C2R4에서 LongCat-Flash-Lite는 HumanEval 71.34%, LiveCodeBench 49.00%를 유지했고, GLM-4.7-Flash는 평균 51.43%로 베이스라인보다 6.22점 높았지만 과제 간 변동이 더 컸다. 4K~32K 검색(RULER S-NIAH)에서는 C4R4가 S-NIAH-1 99.95%, S-NIAH-2 99.60%, S-NIAH-3 87.20%를 기록해 S-NIAH-3에서 최강 양자화 베이스라인을 7.40%p 앞섰고 모든 길이에서 우위였다. 절제 실험에서는 RTN에서 출발해 콘텐츠 변환이 1.22%p, RoPE 변환이 3.19%p를 더해 보호 없는 구성이 51.90%에서 56.31%로 올랐고(SmoothQuant†보다 2.29점, QuaRot†보다 0.65점 높음), 여기에 혼합정밀도 캐시 정책이 1.64%p를 더해 최종 57.95%로 BF16의 57.90%와 대등해졌다.
시스템 효율도 함께 보고된다. 혼합정밀도 정책 기준 C4R4는 4K에서 3.22배, 128K에서 3.59배의 상주 KV 캐시 압축을 달성한다(BF16 버퍼, 중복 패킹 복사본, 페이지 매핑 포함 회계). 1K~1M 길이, 배치 8과 32에서 어텐션 지연은 BF16 대비 -7.42%에서 +4.29% 범위이고 64K 이상에서는 차이가 0.5% 미만이다. 8개 GPU에서 128K 입력·1K 출력 요청 5개를 동시에 처리하는 캐시 압박 워크로드에서 BF16은 32회 선점이 발생한 반면 C4R4는 0회였고, 전체 작업 출력 처리량은 2.61토큰/s에서 13.51토큰/s로 5.168배 증가했다(프리필과 재계산 포함).
개발자 관점에서 이 논문의 실용적 메시지는 두 가지다. 첫째, MLA 캐시를 압축할 때 콘텐츠와 RoPE 키를 같은 기준으로 취급하면 안 된다. 캐시 재구성 오차(NMSE)만 보고 비트를 배분하면 RoPE 경로에서 훨씬 큰 출력 왜곡을 얻는다. 둘째, 경로별 변환을 오프라인에서 가중치에 융합하는 방식이므로 추론 시 추가 변환 연산이 없고, 기존 FlashMLA 디코드 경로에 얹을 수 있는 커널 설계가 함께 제공된다. 다만 sink 4토큰과 최근 128토큰을 BF16으로 보호하는 정책, 그룹 크기 64의 토큰별 비대칭 양자화기, 300스텝·128시퀀스 보정 같은 구체적 설정에 성능이 묶여 있으므로, 자신의 모델·정밀도·서빙 구성에서 재보정과 재검증이 필요하다.
저자들이 명시한 전제와 한계도 분명하다. 윤리 성명에서 낮은 비트 캐시는 집계 벤치마크 정확도가 유지되더라도 모델 동작을 바꿀 수 있으므로, 정확한 모델·정밀도·서빙 구성에서 안전성, 캘리브레이션, 실패 모드를 다시 평가해야 한다고 밝힌다. 5.168배 처리량 이득은 평가한 캐시 압박 워크로드에 특정된 결과라고 못박고, GLM-4.7-Flash에서는 C2R4의 극단적 콘텐츠 압축에 대한 민감도가 커 과제 간 변동이 크다고 서술한다. 또한 이 정밀도에서 직접 비교할 선행 연구가 없어 RTN과 변형한 SmoothQuant†, QuaRot†를 상대로 평가했다는 점도 전제로 깔려 있다.