기술 요약: 양자화가 무료일 때

문제 제기

장문맥 트랜스포머 추론은 일반적으로 키-밸류 (KV) 캐시의 메모리 풋프린트와 대역폭 요구 사항에 의해 제약받습니다. 문맥 길이가 증가함에 따라 (예: 128K 토큰), fp16으로 누적된 KV 텐서는 모델 가중치 자체의 크기를 초과할 수 있으며, 이로 인해 메모리 대역폭이 연산 능력보다 디코드 속도를 제한하는 병목 현상이 발생합니다.

표준 업계 관행은 "회전 후 양자화 (rotate-then-quantize)" 전략 (예: QuaRot, SpinQuant) 을 통해 이를 해결합니다. 이는 가우스 분포를 따르는 활성화 값을 양자화하기 전에 무작위 직교 회전을 적용하여 저비트 (int4/int8) 로 변환하는 방식입니다. 그러나 이 접근법은 전통적으로 품질–지연 시간 트레이드오프로 간주됩니다. 회전 및 역회전의 계산 오버헤드와 양자화 오차가 결합되어 메모리가 절약되더라도 디코드 지연 시간이 일반적으로 5~20% 증가합니다.

저자들은 이 트레이드오프가 **애플 실리콘 (통합 메모리 아키텍처)**에서는 반전된다고 주장합니다. 커널이 매우 최적화되고 퓨전 (fused) 되어 있다면, 3 배 압축된 캐시에서 얻는 대역폭 절감 효과가 양자화 커널의 계산 비용을 상쇄할 수 있다고 가정합니다.

방법론

1. 부호 무작위 푸리에 변환 (SRFT)

이 논문은 회전 메커니즘으로 SRFT(Sign-Randomized FFT) 를 채택하며, $SRFT(x) = \text{pack}(F \cdot \text{diag}(s) \cdot x)$로 정의됩니다. 여기서 $F$는 유니타리 DFT 이고 $s$는 고정된 무작위 부호 벡터입니다.

• 근거: 하다마르 변환 (SRHT) 이 일반적이지만, SRFT 는 2 의 거듭제곱이 아닌 차원에 대한 유연성을 제공하며 16×16 행렬 곱셈 연산 (AMX, 텐서 코어) 과 정렬됩니다.
• 발견: 저자들은 SmolLM2, Qwen2.5, Gemma-3 등 다섯 가지 모델과 다양한 비트 예산에서 SRFT 와 SRHT 가 퍼플렉시티 (PPL) 품질 측면에서 통계적으로 구별할 수 없음을 입증했습니다.

2. 퓨전된 Metal 커널

디스패치 오버헤드를 제거하기 위해 저자는 전체 파이프라인을 수행하는 단일 퓨전된 Metal 계산 커널을 구현했습니다:

1. 로드 및 부호 곱셈: 레지스터 내에서 무작위 부호 대각 행렬을 적용합니다.
2. FFT: log2(d/2) 개의 기수 2 스톡햄 버터플라이 패스를 실행합니다.
3. 패킹: 에르미트 대칭성을 실수 표현으로 언패킹합니다.
4. 스케일링 및 양자화: 채널별 $\lambda$ 재스케일링 (학습 또는 보정), 그룹별 절대 최대값 축소 수행 후 부호 있는 int4/int8 로 반올림합니다.
5. 니블 패킹: SIMD 셔플을 사용하여 두 개의 int4 값을 단일 바이트로 패킹합니다.

이 커널은 torch.mps.compile을 통해 노출되며 HuggingFace Cache 서브클래스 (SRFTInt4Cache) 로 통합됩니다.

3. 고급 스케일링 전략

• 채널별 + 그룹별: 작은 헤드 차원을 가진 모델 (예: Qwen2.5-1.5B, $d=128$) 의 경우, 토큰별 스케일링은 아웃라이어 좌표로 인해 치명적으로 실패합니다. 저자들은 채널별 스케일링(좌표당 하나의 스케일) 과 그룹별 스케일링(16 개 그룹) 을 결합하여 품질을 회복했습니다.
• 학습된 회전: 논문은 학습 가능한 회전 $R$과 스케일 $\lambda$를 추가하는 사후 학습 보정을 조사했습니다. 주요 발견은 SRFT 베이스 없이 $R+\lambda$를 학습할 때 보정 MSE 를 훨씬 더 크게 줄인다는 것 (84.9% 대 50.3%) 이지만, 이는 더 나쁜 다운스트림 PPL을 초래한다는 점입니다. 이는 고정된 무작위 SRFT 가 정규화 역할을 한다는 것을 시사합니다.
• 하우스홀더 회전: $d=256$에서 $k=d/2$ 반사기를 가진 하우스홀더 회전은 전체 카일리 회전의 절반 파라미터로 실질적으로 손실 없는 양자화 ($\Delta PPL \approx 0.000$) 를 달성합니다.

4. CPU 경로 (AMX)

상호 보완적인 CPU 경로는 cblas sgemm를 통해 애플의 AMX 행렬 보조 프로세서를 활용합니다. 길이 8 의 DFT 를 16×16 실수 행렬 곱셈으로 배치함으로써, 저자들은 단일 코어 AMX 피크의 최대 91% 를 달성하여 GPU 디스패치 오버헤드가 지배적인 소규모 배치 환경에서 저지연 대안을 제공합니다.

주요 결과

성능 (지연 시간)

애플 M1(16GB 통합 메모리) 에서 퓨전된 int4 KV 캐시는 엔드투엔드 디코드 지연 시간 측면에서 fp16 을 능가합니다:

• Gemma-3 1B: 2564096 토큰 접두사 전반에 걸쳐 **38% 더 빠른** 속도로 실행됩니다 (부정적 지연 시간 오버헤드).
• Qwen2.5-1.5B: 짧은 문맥 (≤1K 토큰) 에서 0.7~2.6% 더 빠른 속도로 실행됩니다.
• 메커니즘: 커널 오버헤드 (~25 ns/vec) 는 디코드 단계당 전송되는 데이터가 3 배 줄어들어 달성된 메모리 대역폭 절감 효과보다 낮습니다.

품질 (퍼플렉시티)

• SRFT 대 SRHT: 모든 테스트된 모델에서 $\Delta PPL$에 통계적으로 유의미한 차이가 없습니다.
• Qwen2.5-1.5B: 채널별 + 그룹별 스케일링을 적용한 퓨전 커널은 4 비트 $\Delta PPL$을 치명적인 +7975(토큰별 기준) 에서 +638.6(12.5 배 감소) 로 줄여 Python 참조 구현과 일치시킵니다.
• Gemma-3 1B: 4 비트 $\Delta PPL$이 토큰별 기준 +18.7 에서 퓨전된 스케일드 커널을 사용하여 +3.6으로 개선되었습니다.
• 손실 없는 영역: $d=256$에서 학습된 하우스홀더 회전은 $\Delta PPL = 0.000$을 달성합니다.

메모리 효율성

• 압축: 풀 어텐션 레이어에서 3.0 배에서 3.56 배의 지속적 메모리 압축을 달성합니다.
• 실제 영향: Llama-3-70B 에서 128K 문맥의 경우, int4 캐시는 메모리를 ~41 GB(fp16) 에서 ~10.9 GB 로 줄입니다.

중요성 및 주장

이 논문은 통합 메모리 하드웨어에서 **양자화가 "무료"(처리량 긍정적)**라고 주장하며, 저비트 양자화가 항상 지연 시간 패널티를 수반한다는 표준 업계 가정을 반전시킵니다.

1. 트레이드오프의 반전: 주요 기여는 애플 실리콘에서 3 배 압축된 캐시에서 얻는 대역폭 절감 효과가 양자화 커널의 계산 비용을 초과한다는 것을 입증하는 것입니다. 이는 노트북 및 엣지 장치에서 지연 시간에 민감한 장문맥 추론을 위해 int4 KV 캐시가 fp16 보다 엄격하게 우월하게 만듭니다.
2. 하드웨어 인식 설계: 이 연구는 회전 (FFT 대 하다마르) 의 선택이 수학적 우월성보다는 하드웨어 정렬 (예: AMX 16×16 타일과 일치하는 FFT 의 기수 8 구조) 에 더 관련이 있음을 강조합니다.
3. 보정 대 PPL 분리: 논문은 보정 MSE 를 최소화하는 것이 더 나은 다운스트림 PPL 을 보장하지 않는 새로운 현상을 분리해냈으며, 고정된 무작위 회전이 필요한 정규화 사전 역할을 한다는 것을 시사합니다.
4. 배포 가능성: 전체 파이프라인을 단일 Metal 커널로 퓨전하고 HuggingFace 와 통합함으로써, 저자들은 하드웨어 아키텍처 (통합 메모리) 와 커널 최적화 (퓨전) 가 정렬되어 있다면 모델 품질이나 속도를 희생하지 않고 공격적인 4 비트 양자화를 배포할 수 있음을 증명합니다.

저자들은 표준 양자화 트레이드오프가 별도의 HBM 과 높은 디스패치 오버헤드를 가진 GPU 에 대해서는 유효하지만, 노트북 및 엣지 추론이 작동하는 통합 메모리 영역에서는 반전된다고 결론지었습니다.