기술 요약: OSCAR - 2 비트 KV 캐시 양자화를 위한 오프라인 스펙트럴 공분산 인지 회전

1. 문제 제기

대규모 언어 모델 (LLM) 의 장문맥 추론은 increasingly Key-Value(KV) 캐시의 메모리 풋프린트에 의해 병목 현상을 겪고 있습니다. 문맥 길이, 배치 크기, 모델 깊이가 증가함에 따라 캐시는 상당한 GPU 메모리를 소모하여 처리량과 확장성을 제한합니다. KV 캐시를 초저정밀도 (예: INT2) 로 압축하면 8 배의 메모리 감소를 제공하지만, 치명적인 정확도 문제에 직면합니다: KV 활성화 값에는 심각한 채널별 이상치 (outliers) 가 존재합니다.

저비트 양자화 하에서 이러한 이상치는 양자화 스케일을 지배하여, 대부분의 '정상' 값을 소수의 유효 양자화 레벨로 강제 배치하고 주의 (attention) 품질을 심각하게 저하시킵니다. 기존 회전 기반 방법들 (예: Hadamard 변환) 은 채널 간 에너지를 재분배하여 이러한 이상치를 완화하려 시도합니다. 그러나 이러한 방법들은 **데이터 무관 (data-oblivious)**하며, 원시 캐시 분포를 완화하지만 하류 주의 메커니즘이 실제로 소비하는 특정 방향과 정렬되지 못합니다. 단 4 개의 양자화 레벨만 존재하는 INT2 에서, 이러한 불일치는 단순한 회전 방법들을 붕괴시켜 추론 작업에서 거의 0 에 가까운 정확도를 초래합니다.

2. 방법론: OSCAR

저자들은 OSCAR(Offline Spectral Covariance-Aware Rotation) 을 제안합니다. 이는 오프라인으로 주의 인지 공분산 구조를 추정하여 고정된 회전 및 클리핑 임계값을 도출하는 방법입니다. 기존 연구가 원시 유클리드 재구성 오차 ($\|K - \hat{K}\|_F$) 를 최소화하는 것과 달리, OSCAR 은 하류 주의 점수 및 출력의 왜곡을 최소화합니다.

2.1 이론적 동기

이 논문은 최적의 회전 목표가 주의 계산 자체의 통계에서 도출되어야 한다고 주장합니다:

• 키 (Keys, $K$) 의 경우: 주의 로그이트 (logits) 의 왜곡은 키 공분산 $K^\top K$뿐만 아니라 쿼리 공분산 $Q^\top Q$에 의해 제어됩니다.
• 값 (Values, $V$) 의 경우: 주의 출력의 왜곡은 주의 점수 $S$로 가중되며, 이는 $V^\top S^\top S V$를 목표 공분산으로 이끕니다.

2.2 알고리즘 설계

OSCAR 은 두 단계로 운영됩니다:

A. 오프라인 보정 단계:

1. 공분산 추정: 소규모 보정 데이터셋을 사용하여 경험적 목표 공분산을 계산합니다:
   • $C_Q = \frac{1}{N} \sum q_n^\top q_n$ (쿼리 인지 키 목표).
   • $C_S = \frac{1}{N} V^\top S^\top S V$ (점수 인지 값 목표).
2. 회전 유도: 이 방법들은 이러한 행렬들에 대해 고유분해 (eigen-decomposition) 를 수행하여 최적의 고유기저 ($U_Q$ 및 $U_S$) 를 찾습니다.
3. 복합 회전 구성: 최종 회전 행렬 ($R_K, R_V$) 은 세 가지 구성 요소의 곱으로 구성됩니다:
   • $U_Q/U_S$: 채널을 중요도 방향과 정렬시키는 주의 인지 고유기저.
   • $H_{Had}$: 채널 간 중요도 지표를 균등화하고 이상치 에너지를 분산시키는 Hadamard 변환.
   • $P_{br}$: 양자화 그룹 내에서 균형 잡힌 동적 범위를 보장하기 위해 고분산 및 저분산 채널을 교차 배치하는 비트 반전 순열.
   • 공식: $R_K = U_Q H_{Had} P_{br}$ 및 $R_V = U_S H_{Had} P_{br}$.
4. 클리핑 임계값: 회전된 보정 활성화 값에서 레이어별 클리핑 임계값을 추정하여 이상치를 제어합니다.

B. 온라인 서비스 단계:

• 혼합 정밀도 레이아웃: 시스템은 '싱커 (sink)' 토큰 (초기) 과 '최근' 토큰을 전체 BF16 정밀도로 유지하는 페이지드 KV 캐시를 유지합니다. 역사적 중간 문맥은 고정된 OSCAR 회전 및 클리핑을 적용한 후 INT2 로 저장됩니다.
• 커널 통합: 저자들은 SGLang 서비스 프레임워크를 위한 커스텀 퓨즈드 Triton 커널을 개발했습니다. 이 커널은 INT2 히스토리에 대한 언패킹, 역양자화, 역회전 및 주의 계산을 처리하여 페이지드 주의 및 접두사 캐싱과 호환성을 보장하면서 추가 메모리 트래픽을 도입하지 않습니다.

3. 주요 기여

1. 누락된 목표의 식별: 이 논문은 일반적인 회전들이 이상치를 분산시키지만 INT2 에서 실패하는 이유는 하류 주의와 정렬되지 않기 때문임을 식별합니다. 회전 목표는 원시 캐시 재구성이 아닌 주의 통계 ($Q^\top Q$ 및 $V^\top S^\top S V$) 에 의해 유도되어야 합니다.
2. OSCAR 프레임워크: 오프라인 보정을 사용하여 주의 인지 회전을 도출하는 완전히 배포 가능한 INT2 KV 캐시 양자화 시스템입니다. 이론적 분석은 이러한 목표가 고정된 오차 대리 가정 하에서 최적임을 보여줍니다.
3. 생산 준비 시스템: 페이지드 KV 캐시 레이아웃 및 접두사 캐싱을 지원하는 SGLang 에 통합된 커스텀 INT2 주의 커널입니다. 이는 BF16 및 INT4 기반 모델들이 메모리 부족에 직면하는 단일 H100 에서 10 만 토큰 입력으로 최대 28 개의 동시 요청을 처리할 수 있음을 보여주며, 현대 서비스 프레임워크 (SGLang, vLLM) 와의 원활한 통합을 달성하면서 높은 처리량을 유지합니다.

4. 실험 결과

저자들은 추론 모델 (Qwen3-4B, Qwen3-8B, Qwen3-32B) 과 최첨단 모델 (GLM-4.7, 358B 파라미터) 에 대해 다섯 가지 벤치마크 (AIME25, GPQA, HumanEval, LiveCodeBench, MATH500) 및 장문맥 검색 (RULER-NIAH 최대 128k) 에서 OSCAR 을 평가했습니다.

• 정확도 보존:

  • Qwen3-4B-Thinking 및 Qwen3-8B에서 OSCAR 은 BF16 대비 정확도 격차를 각각 3.78 및 1.42 포인트로 줄입니다. 반면, 단순 회전 INT2 는 거의 0 에 가까운 정확도로 붕괴됩니다.
  • Qwen3-32B 및 GLM-4.7에서 OSCAR 은 BF16 과 실질적으로 동등한 성능을 유지합니다 (0.02~0.27 포인트 이내),而其他 INT2 기준 모델들은 심각한 저하를 겪습니다.
  • 장문맥 검색 (RULER-NIAH) 에서 OSCAR 은 128k 토큰까지 견고한 성능을 유지하는 반면, 단순 INT2 및 Hadamard 전용 방법은 훨씬 짧은 길이 (예: 16k) 에서 붕괴됩니다.

• 시스템 성능:

  • 메모리: BF16 대비 KV 캐시 메모리 풋프린트를 약 8 배 줄입니다.
  • 처리량: 동일한 메모리 예산 하에서 대규모 배치 크기에서 처리량을 최대 7 배 향상시킵니다.
  • 지연 시간: 메모리 대역폭 오버헤드 감소로 인해 BF16 대비 배치 크기 1 디코딩 속도를 최대 3 배 가속화합니다.
  • 확장성: 단일 H100 에서 10 만 토큰 입력으로 최대 28 개의 동시 요청을 처리할 수 있음을 보여주며, 이는 BF16 및 INT4 기준 모델들이 메모리 부족에 직면하는 영역입니다.

5. 중요성 및 주장

이 논문은 OSCAR 이 초저비트 (INT2) KV 캐싱의 이론적 약속과 실제 배포 가능성 사이의 격차를 성공적으로 해소했다고 주장합니다. 회전 목표를 원시 데이터 재구성에서 주의 인지 공분산으로 전환함으로써, OSCAR 은 단순 양자화로 인해 손실된 정확도를 회복하면서 2 비트 저장의 막대한 메모리 및 처리량 이점을 유지합니다.

저자들은 그들의 접근 방식이 단순한 알고리즘 개선이 아닌 시스템 수준의 솔루션임을 강조합니다. 페이지드 주의 및 접두사 캐싱과 통합되는 커스텀 커널을 제공함으로써, 그들은 INT2 양자화가 현대 LLM 서비스에 의존하는 효율성 기능 (예: 접두사 재사용) 을 희생하지 않고도 (SGLang 과 같은) 생산 서비스 프레임워크에서 사용될 수 있음을 보여줍니다. 이 연구는 장문맥 추론을 위해 성공적인 양자화의 열쇠는 압축 전략을 주의 메커니즘의 특정 통계적 요구 사항과 정렬시키는 데 있음을 시사합니다.