[onnxruntime] ONNX Runtime CPU QMoE 최적화: INT2 AVX-512/AVX2 커널 도입 및 수치 안정성 개선

PR 링크: microsoft/onnxruntime#33785 상태: Merged | 변경: +1971 / -869

들어가며

최근 대규모 언어 모델(LLM)의 효율적인 추론을 위해 Mixture of Experts(MoE) 구조가 널리 사용되고 있습니다. 특히, 가중치를 양자화하여 메모리 사용량을 줄이는 QMoE(Quantized MoE)는 필수적인 기술입니다. 이번 Microsoft ONNX Runtime의 업데이트는 CPU 환경에서 INT2 가중치를 직접 처리할 수 있는 고성능 커널을 도입하고, 기존 CUDA 구현체와의 수치적 일관성을 맞추는 데 초점을 맞추고 있습니다.

코드 분석

1. 고성능 INT2 SIMD 커널 도입

cmake/onnxruntime_mlas.cmake를 통해 AVX-512 및 AVX2용 INT2 커널이 새롭게 추가되었습니다. 이는 기존의 INT4/INT3 커널과 동일한 아키텍처를 공유하며, 런타임 디스패치를 통해 호스트 CPU의 기능을 최대한 활용합니다.

+    "${ONNXRUNTIME_ROOT}/contrib_ops/cpu/moe/qmoe_shared_layout_int2_avx512.cc"
+    "${ONNXRUNTIME_ROOT}/contrib_ops/cpu/moe/qmoe_shared_layout_int2_avx2.cc"

2. 수치 안정성: 오버플로우 방지 및 반올림 정책

가장 중요한 변경 사항 중 하나는 수치적 안정성입니다. 특히 BF16 스케일이 매우 클 경우 발생하는 오버플로우 문제를 해결하기 위해 스케일을 반으로 나누고 블록 합을 두 배로 늘리는 기법을 도입했습니다.

// Before (Scalar path)
return scale_is_fp32 ? scales_f32[s_idx] : ...;

// After (Scalar path with overflow protection)
constexpr float kHalveThreshold = 0x1p127f;
return std::fabs(scale) >= kHalveThreshold ? ScaleTerm{0.5f * scale, 2.0f} : ScaleTerm{scale, 1.0f};

또한, FC2 레이어의 bias 연산 시 CUDA 구현체와 일관성을 맞추기 위해, bias를 더하기 전에 먼저 projection 값을 반올림(rounding)하도록 변경했습니다.

// After: Rounding before bias addition to match CUDA SIMT/Tensor Core
float value = RoundElement(fc2[r * hidden + n], boundary_type);
if (view.fc2_bias) value += read_fc2_bias(n);
destination[n] = value;

왜 이게 좋은가

1. 성능 향상: 새로운 AVX-512/AVX2 커널은 기존 scalar 구현 대비 훨씬 높은 처리량을 제공합니다. 실제 벤치마크 결과, BF16 block-32, M=1 decode GEMM 연산에서 오버플로우 처리를 포함하고도 34.3us 수준의 매우 낮은 지연 시간을 기록했습니다.
2. 수치적 일관성: CUDA(SM80) 구현체와 동일한 rounding 정책을 적용함으로써, CPU와 GPU 간의 추론 결과 차이(parity)를 제거했습니다. 이는 모델 배포 시 플랫폼 간 결과가 달라지는 문제를 방지합니다.
3. 견고한 예외 처리: SiLU 활성화 함수에서 발생할 수 있는 -inf gate 문제를 해결하여 NaN 발생을 방지하고, 0으로 수렴하도록 보장하여 모델의 안정성을 높였습니다.

교훈

이번 최적화는 단순히 속도를 높이는 것뿐만 아니라, '수치적 정확성(Numerical Correctness)'이 고성능 연산 라이브러리에서 얼마나 중요한지 보여줍니다. 특히 서로 다른 하드웨어 가속기(CPU vs GPU) 간의 연산 순서와 반올림 정책을 맞추는 것은 복잡한 LLM 추론 엔진 개발에서 필수적인 과정입니다.

참고 자료

• ONNX Runtime MLAS
• AVX-512 Intrinsics Guide

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.