CubicQuant: Parametric Non-Uniform Codebooks for High-Throughput LLM Inference with 1-8-Bit Weights

기술 요약: CubicQuant

문제 정의

대규모 언어 모델(LLM) 추론은 두 가지 상충하는 목표 사이의 균형을 맞춰야 합니다: 양자화된 가중치의 재구성 레벨(reconstruction levels)은 모델 가중치의 다양한 국소 통계(종종 가우시안 또는 라플라스와 같은 비균등 분포를 따름)에 적응해야 하지만, 표현 방식은 고처리량 GPU 하드웨어에서 직접 실행될 수 있을 만큼 조밀하고 규칙적이어야 합니다.

기존의 스칼라 양자화 방법들은 이 트레이드오프 스펙트럼의 서로 다른 지점에 위치합니다:

• 균등 정수 양자화 (Uniform Integer Quantization): 밀집한 패킹과 단순한 산술 연산을 제공하지만, 모든 그룹을 선형 그리드에 제한하여 비균등 가중치 분포를 효율적으로 포착하지 못합니다.
• 저비트 부동 소수점 (예: FP8): 고정된 지수-가수 구조를 통해 레벨을 재배치하여 미리 결정된 동적 범위를 선호하지만, 특정 가중치 통계에 대한 유연성이 부족합니다.
• 학습된 코드북 (Learned Codebooks): 모든 레벨을 학습함으로써 유연성을 얻지만, 불규칙한 디코딩, 복잡한 메타데이터 구조, 그리고 표준 GPU 행렬 곱 연산 경로와의 낮은 정렬 문제를 겪습니다.

과제는 효율적인 GPU 실행을 위한 정수 코드의 규칙성을 유지하면서도, 재구성 레벨이 국소 가중치 통계에 매개변수적으로 적응할 수 있는 형식을 만드는 것입니다.

방법론: CubicQuant

CubicQuant는 밀집한 정수 코드 스트림을 보존하면서 각 가중치 그룹 내에서 독립적으로 재구성 레벨을 적응시키는 **매개변수적 비균등 스칼라 형식(parametric non-uniform scalar format)**을 도입합니다.

1. 표현 (Representation)

• 페이로드 (Payload): 가중치는 $-M$에서 $M$ 사이의 범위($M = 2^{B-1}-1$, 여기서 $B$는 비트 폭)를 갖는 부호 있는 정수 코드($k$)로 저장됩니다. 이 형식은 1~8비트 페이로드를 지원합니다.
• 매개변수적 레벨 함수 (Parametric Level Function): 선형 그리드 대신, 재구성 레벨은 **단조 삼차 곡선(monotonic cubic curve)**에 의해 결정됩니다. 정규화된 크기 $t = |k|/M$에 대하여, 재구성 함수는 다음과 같습니다:
  $$q(t) = t [a + t (b + ct)]$$
  여기서 $c = 1 - a - b$입니다.
  • 매개변수: 각 그룹은 두 개의 무차원 형상 계수($a, b$)와 하나의 FP32 스케일($s$)을 저장합니다.
  • 제약 조건: 곡선은 단조성($q'(t) > 0$)을 유지하며 $q(0)=0$ 및 $q(1)=1$을 만족하도록 제약됩니다. 이는 정확한 제로(zero)와 정확한 끝점($\pm s$)을 보장합니다.
  • 특수 사례: $a=1, b=0$인 경우, 균등 대칭 정수 양자화는 이 계열의 정확한 특수 사례가 됩니다.

2. 메타데이터 및 저장 (Metadata and Storage)

• 그룹화 (Grouping): 가중치는 크기 $G$의 그룹으로 나뉩니다.
• 메타데이터: 각 그룹은 하나의 FP32 스케일과 두 개의 FP16 형상 계수를 저장합니다.
• 유효 폭 (Effective Width): 총 저장 용량은 가중치당 $B + 64/G$ 비트입니다. 이는 모델 수준의 평균이나 복잡한 룩업 테이블에 숨겨진 메타데이터와 대조되는 명시적인 메타데이터 비용입니다.

3. 피팅 목적 함수 (Fitting Objectives)

본 논문은 두 가지 주요 피팅 전략을 정의합니다:

• 연속 참조 목적 함수 (Continuous Reference Objective): 원래 가중치와 연속적인 Cubic 재구성 레벨 사이의 평균 제곱 오차(MSE)를 최소화합니다. 이는 데이터가 필요 없는(data-free), 가중치 전용 추정 방식입니다.
• Dynamic-A8 캐리어 인식 목적 함수 (Dynamic-A8 Carrier-Aware Objective): 동적으로 양자화된 INT8 활성화 함수를 지원하기 위해, 피팅 목적 함수는 연속 재구성 및 이를 부호 있는 INT8 캐리어 그리드($r(k) = \text{round}(127 \cdot \text{sign}(k) \cdot q(|k|/M))$로의 투영)의 결합 오차를 최소화하도록 수정됩니다. 이를 통해 단일 직렬화된 체크포인트가 모델-dtype 및 Dynamic-A8 실행 경로 모두를 지원할 수 있습니다.

4. GPU 실행 (GPU Execution)

CubicQuant는 전체 정밀도 가중치 텐서를 실체화하는 것을 피합니다. 대신 융합 디코드-및-컴퓨트(fused decode-and-compute) 방식을 채택합니다:

• 타일 로컬 재구성 (Tile-Local Reconstruction): GPU 커널은 패킹된 페이로드와 그룹 메타데이터를 로드하고, 레벨을 재구성(룩업 테이블 또는 Horner-form 평가를 통해)한 후, 즉시 도트 곱 루프에서 소비합니다.
• 이중 경로 (Dual Paths):
  • Model-Dtype 경로: FP16/BF16 활성화를 위해 연속적인 값을 재구성합니다.
  • Dynamic-A8 경로: 재구성된 레벨을 INT8 캐리어로 매핑하여 정수 도트 곱을 수행하고, INT32에서 누적한 후 스케일링합니다.
• 커널 선택 (Kernel Selection): 구현체는 행 수와 비트 폭에 따라 최적의 성능을 내기 위해 형상 및 장치 의존적 커널 선택(예: 좁은 GEMV 대 밀집 타일링)을 사용합니다.

주요 기여 (Key Contributions)

1. 매개변수적 코드북 계열: 균등 부호 있는 정수 양자화를 정확한 특수 사례로 포함하는 대칭적, 그룹별, 두 매개변수 기반의 스칼라 비균등 코드북 계열을 제시합니다.
2. 패킹된 형식: FP32 스케일링과 닫힌 형태의 저장 오버헤드($B + 64/G$ bits/weight)를 갖는 1~8비트 가중치 스팬 체크포인트 형식을 제공합니다.
3. 피팅 절차: 그룹별 참조 피팅 절차와 Dynamic-A8-캐리어 인식 목적 함수를 제공하며, 이는 더 풍부한 데이터 기반 오차 추정기와 분리 가능합니다.
4. 직접 GPU 실행: 모델-dtype 및 Dynamic-A8 활성화를 위한 패킹된 가중치 GPU 커널과 장치 인식 캘리브레이션을 구현했습니다.
5. 분석적 및 경험적 근거: 균등, 가우시안, 라플라스 분포에 대한 인구 통계적 왜곡(population distortion) 유도와 함께 H200 커널 측정치를 제시합니다.

결과 (Results)

이론적 왜곡 분석 (Theoretical Distortion Analysis)

• 균등 분포 (Uniform Distribution): 선형 그리드(Uniform INT)가 인구 통계적으로 최적입니다. CubicQuant는 이 성능을 일치시켜, 균등한 데이터에 대해 성능을 저하시키지 않음을 확인했습니다.
• 가우시안 및 라플라스 분포 (Gaussian and Laplace Distributions):
  • CubicQuant는 최적으로 클리핑된 균등 정수 양자화보다 성능이 뛰어납니다.
  • 가우시안 분포(W4–W8)에 대해, 균등 정수와 최적의 자유 Lloyd–Max 코드북 사이의 격차 중 **84.8%~96.1%**를 회복합니다.
  • 라플라스 분포에 대해서는 **91.4%~96.9%**의 격차를 회복합니다.
  • 모든 테스트된 비트 폭에서 최적의 열거된 유한 부동 소수점 형식(예: E2M5)보다 우수한 성능을 보입니다.

유한 그룹 실험 (Finite-Group Experiments)

$G=128$, 15,360개 샘플을 사용한 실험 결과:

• W4 성능: 최적으로 클리핑된 4비트 균등 정수 대비, CubicQuant는 재구성 RMSE를 다음과 같이 감소시켰습니다:
  • 균등 샘플: 3.90%
  • 가우시안 샘플: 13.49%
  • 라플라스 샘플: 28.14%
• FP와의 비교: 최적의 열거된 4비트 부동 소수점 형식 대비, 감소율은 균등(3.90%), 가우시안(9.44%), 라플라스(6.27%)였습니다.
• 메커니즘: 균등 데이터에서의 이득은 인구 통계적 최적값이 선형임에도 불구하고, 그룹 내의 유한 샘플 변동에 적응함으로써 발생합니다.

커널 성능 (H200)

• 워크로드 의존적 교차 지점 (Workload-Dependent Crossover): H200 GPU에서의 예비 측정 결과, Model-Dtype(A16)과 Dynamic-A8 실행 간의 교차 지점이 발견되었습니다.
  • 좁은 GEMV(작은 행 수, $M=1$)의 경우, 설정 비용이 낮아 Model-Dtype 실행이 더 빠릅니다.
  • 행 수가 증가하면($M \ge 16$), Dynamic-A8이 유리해지며, 비트 폭과 형상에 따라 1.6배에서 4.4배의 속도 향상을 보입니다.
• 부정적 결과 (Negative Result): "지속적 활성화 캐리어(persistent activation carrier)" 설계(연산자 사이에 INT8 활성화를 유지하는 방식)는 동기화 및 리덕션 오버헤드가 대역폭 절감 효과를 초과하여 성능 개선에 실패했습니다.

의의 및 주장 (Significance and Claims)

본 논문은 CubicQuant 형식의 표현 가능성과 실행 가능성을 확립한다고 주장합니다.

• 복잡성 없는 적응성: CubicQuant는 압축된 매개변수 곡선이 효율적인 GPU 실행에 필요한 정규 스칼라 비트스트림을 포기하지 않으면서도, 자유 코드북의 적응성 상당 부분을 회복할 수 있음을 보여줍니다.
• 이중 경로 호환성: 이 형식은 단일 직렬화된 표현으로부터 연속적인 모델-dtype 실행과 정수 기반의 Dynamic-A8 실행을 모두 성공적으로 지원합니다.
• 제한된 범위: 저자는 본 논문이 다음을 주장하지 않음을 명시합니다:
  • 다운스트림 모델 품질(퍼플렉시티, 추론 능력 등)에 대한 측정은 이루어지지 않았습니다.
  • 보편적인 엔드 투 엔드 속도 향상을 의미하지 않습니다; 성능은 워크로드에 따라 다릅니다.
  • 모든 GPU 세대에 대한 일반적 적용성을 의미하지 않습니다 (증거는 Hopper/H200에 국한됨).
  • 데이터 프리(data-free) 참조 추정기가 활성화 인식 또는 2차 방법보다 우월하다는 것을 주장하지 않습니다; 이 형식은 그러한 방법들과 호환되도록 설계되었습니다.

요약하자면, CubicQuant는 고정된 그리드와 자유 코드북 사이의 중간 지점을 제공하며, 현대 GPU 아키텍처에서 직접 실행 가능한 수학적으로 근거 있는 매개변수적 접근 방식을 제공합니다.