[onnxruntime] ONNX Runtime, x86 CPU에서 FP16 LayerNorm 및 RMSNorm 성능 최적화: AVX2 활용
PR 링크: microsoft/onnxruntime#32715 상태: Merged | 변경: +400 / -1
들어가며
최근 Microsoft의 ONNX Runtime 레포지토리에서는 x86 CPU 아키텍처에서 FP16(반정밀도 부동소수점) 데이터 타입을 사용하는 Layer Normalization 및 RMS Normalization 연산의 성능을 최적화하는 중요한 변경이 있었습니다. 특히 트랜스포머와 같은 딥러닝 모델에서 빈번하게 사용되는 이 연산들은 기존에는 스칼라 연산과 개별 요소 변환에 많은 시간을 소요했습니다. 본 PR은 AVX2, FMA, F16C와 같은 최신 CPU 명령어셋을 활용하여 이러한 병목 현상을 해결하고, 특히 긴 시퀀스 길이를 가진 데이터에서 상당한 성능 향상을 이끌어냈습니다.
이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 왜 이러한 최적화가 효과적인지, 그리고 이를 통해 얻을 수 있는 일반적인 교훈은 무엇인지 살펴보겠습니다.
코드 분석
이번 최적화는 주로 ONNX Runtime의 MLAS(Microsoft Linear Algebra Subprograms) 라이브러리 내에서 이루어졌습니다. 핵심 변경 사항은 다음과 같습니다.
1. MLAS 빌드 설정 변경 (cmake/onnxruntime_mlas.cmake)
AVX2 명령어셋을 사용하는 커널 컴파일 시 F16C(FP16 변환) 명령어셋도 함께 활성화하도록 변경되었습니다.
Before:
- set_source_files_properties(${mlas_platform_srcs_avx2} PROPERTIES COMPILE_FLAGS "-mavx2 -mfma")
+ set_source_files_properties(${mlas_platform_srcs_avx2} PROPERTIES COMPILE_FLAGS "-mavx2 -mfma -mf16c")
After:
+ set_source_files_properties(${mlas_platform_srcs_avx2} PROPERTIES COMPILE_FLAGS "-mavx2 -mfma -mf16c")
설명:
기존에는 AVX2와 FMA 명령어만 활성화했지만, FP16 데이터를 효율적으로 처리하기 위해서는 F16C 명령어셋이 필수적입니다. -mf16c 플래그를 추가함으로써 컴파일러는 FP16과 FP32 간의 변환을 위한 최적화된 하드웨어 명령어를 사용할 수 있게 됩니다. 이는 layernorm_kernel_avx2.cpp 파일 내에서 _mm256_cvtph_ps (FP16 to FP32) 및 _mm256_cvtps_ph (FP32 to FP16)와 같은 AVX2/F16C 내장 함수(intrinsics)를 사용하는 코드의 기반을 마련합니다.
리뷰어의 지적([Copilot on onnxruntime/core/mlas/lib/layernorm_kernel_avx2.cpp])처럼, 32비트 환경에서도 이 플래그가 필요하며, 이 변경은 해당 문제를 해결합니다.
2. MLAS 헤더 파일에 FP16 LayerNorm 함수 선언 (onnxruntime/core/mlas/inc/mlas.h)
FP16 입력을 처리하는 새로운 MLAS 함수 MlasLayerNormF16이 선언되었습니다.
추가:
+/**
+ * @brief Compute LayerNorm or RMSNorm (simplified) for one row of IEEE FP16 data.
+ * Scale and bias are supplied as float32 to match the CPU LayerNorm
+ * prepacking path. Returns false when no optimized kernel is available.
+ */
+bool
+MLASCALL
+MlasLayerNormF16(
+ const uint16_t* Input,
+ const float* Scale,
+ const float* Bias,
+ uint16_t* Output,
+ float* MeanOut,
+ float* InvStdDevOut,
+ size_t NormSize,
+ float Epsilon,
+ bool Simplified
+);
설명:
이 함수는 FP16 입력(Input), FP32 스케일(Scale) 및 편향(Bias)을 받아 FP16 출력(Output)을 생성합니다. 또한, 선택적으로 평균(MeanOut)과 표준편차의 역수(InvStdDevOut)를 FP32로 반환할 수 있습니다. Simplified 플래그는 편향(bias)이 없는 경우를 나타냅니다. 이 함수는 내부적으로 최적화된 커널을 호출하며, 커널이 사용 불가능하거나 입력 크기가 작은 경우 false를 반환하여 스칼라 경로로 폴백합니다.
3. MLAS 구현 파일에 FP16 LayerNorm 함수 구현 (onnxruntime/core/mlas/lib/layernorm.cpp)
새롭게 선언된 MlasLayerNormF16 함수의 실제 구현이 추가되었습니다.
추가:
+bool
+ MLASCALL
+ MlasLayerNormF16(
+ const uint16_t* Input,
+ const float* Scale,
+ const float* Bias,
+ uint16_t* Output,
+ float* MeanOut,
+ float* InvStdDevOut,
+ size_t NormSize,
+ float Epsilon,
+ bool Simplified
+ )
+{
+ auto kernel = GetMlasPlatform().LayerNormF16Kernel;
+ if (kernel == nullptr || NormSize < 16) {
+ return false;
+ }
+
+ kernel(Input, Scale, Bias, Output, MeanOut, InvStdDevOut, NormSize, Epsilon, Simplified);
+ return true;
+}
설명:
이 구현은 GetMlasPlatform().LayerNormF16Kernel을 통해 현재 CPU 아키텍처에 맞는 최적화된 커널 포인터를 가져옵니다. 만약 커널이 없거나(kernel == nullptr) 처리할 행의 크기(NormSize)가 16 미만이면, 최적화된 경로를 사용하지 않고 false를 반환합니다. 이는 작은 크기의 입력에 대해서는 벡터 연산의 오버헤드가 스칼라 연산보다 클 수 있기 때문입니다. 16개 이상의 요소로 구성된 행에 대해서는 최적화된 커널이 호출됩니다.
4. AVX2 기반 FP16 LayerNorm 커널 구현 (onnxruntime/core/mlas/lib/layernorm_kernel_avx2.cpp)
이 파일은 AVX2 명령어셋을 사용하여 FP16 LayerNorm 및 RMSNorm을 효율적으로 계산하는 핵심 로직을 담고 있습니다.
주요 변경 및 추가:
+#include "mlas_float16.h"
...
+void MLASCALL
+MlasLayerNormKernelF16Avx2(
+ const uint16_t* Input,
+ const float* Scale,
+ const float* Bias,
+ uint16_t* Output,
+ float* MeanOut,
+ float* InvStdDevOut,
+ size_t NormSize,
+ float Epsilon,
+ bool Simplified)
+{
+ assert(!Simplified || Bias == nullptr);
+ const size_t n = NormSize;
+ float mean_val = 0.0f;
+ float inv_denom;
+
+ if (Simplified) {
+ // ... FP16 to FP32 conversion and FP32 computation ...
+ // Uses __m256_cvtph_ps, _mm256_fmadd_ps, etc.
+ } else {
+ // ... FP16 to FP32 conversion and FP32 computation ...
+ // Uses __m256_cvtph_ps, _mm256_fmadd_ps, etc.
+ }
+
+ // ... FP32 to FP16 conversion and output writing ...
+ // Uses __m256_cvtps_ph
+...
+ if (MeanOut != nullptr) { *MeanOut = mean_val; }
+ if (InvStdDevOut != nullptr) { *InvStdDevOut = inv_denom; }
+}
설명: 이 커널은 다음과 같은 단계를 거쳐 FP16 데이터를 처리합니다:
- FP16 -> FP32 변환: AVX2 레지스터(
__m256)를 사용하여 한 번에 8개의 FP16 입력 요소를 로드하고_mm256_cvtph_ps내장 함수를 통해 FP32로 변환합니다. 이는 기존의 스칼라 방식보다 훨씬 빠릅니다. - FP32 연산: 변환된 FP32 값들을 사용하여 평균(mean), 분산(variance) 또는 합계(sum)를 계산합니다. LayerNorm의 경우, 평균 누적 시
double정밀도를 사용하여 정확도를 높입니다 (_mm256_add_pd). RMSNorm의 경우, 합계(sum of squares)를 계산합니다. - 정규화 및 스케일/편향 적용: 계산된 평균과 표준편차(또는 합계)를 사용하여 정규화(
(x - mean) / inv_stddev)를 수행합니다. 이후 FP32 스케일(Scale)과 편향(Bias)을 적용합니다. 이 과정에서도 FMA(_mm256_fmadd_ps)와 같은 명령어 활용으로 연산 효율을 높입니다. - FP32 -> FP16 변환: 최종 FP32 결과를
_mm256_cvtps_ph내장 함수를 사용하여 다시 FP16으로 변환합니다. 이 때 반올림 모드(_MM_FROUND_TO_NEAREST_INT)가 지정됩니다. 리뷰어xadupre의 지적에 따라 MSVC 빌드 호환성을 위해 반올림 모드 관련 플래그가 수정되었습니다. - 결과 저장: 변환된 FP16 결과는 벡터 레지스터(
__m128i)에 저장된 후 메모리에 기록됩니다.
이 커널은 특히 NormSize >= 16 조건에서만 활성화되어, 벡터 연산의 이점을 최대한 활용합니다.
5. MLAS 플랫폼 및 커널 디스패치 업데이트 (onnxruntime/core/mlas/lib/mlasi.h, onnxruntime/core/mlas/lib/platform.cpp)
MLAS 플랫폼 정보에 FP16 LayerNorm 커널 포인터를 추가하고, CPUID를 통해 F16C 지원 여부를 확인하여 해당 커널을 등록하는 로직이 추가되었습니다.
mlasi.h 변경:
typedef void(MLASCALL MLAS_LAYERNORM_F16_KERNEL)(
...
);
...
MLAS_LAYERNORM_F32_KERNEL* LayerNormF32Kernel{nullptr};
+ MLAS_LAYERNORM_F16_KERNEL* LayerNormF16Kernel{nullptr};
platform.cpp 변경 (일부):
this->LayerNormF32Kernel = &MlasLayerNormKernelAvx2;
+ if ((Cpuid1[2] & (1u << 29)) != 0) { // Check for F16C support
+ this->LayerNormF16Kernel = &MlasLayerNormKernelF16Avx2;
+ }
설명:
mlasi.h에서는 MLAS_LAYERNORM_F16_KERNEL 타입과 LayerNormF16Kernel 포인터를 MlasPlatform 구조체에 추가하여, 플랫폼별로 등록 가능한 커널 목록을 확장했습니다. platform.cpp에서는 CPUID 정보를 확인하여 F16C 명령어셋(Cpuid1[2] & (1u << 29))을 지원하는 경우에만 MlasLayerNormKernelF16Avx2 커널을 LayerNormF16Kernel 포인터에 할당합니다. 이를 통해 런타임 시 CPU의 기능을 동적으로 감지하고 최적화된 커널을 선택적으로 사용할 수 있게 됩니다.
6. CPU Provider의 LayerNorm 구현 업데이트 (onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc)
CPU Provider 레벨에서 FP16 LayerNorm 연산을 처리할 때, MLAS의 새로운 MlasLayerNormF16 함수를 호출하도록 로직이 추가되었습니다.
리뷰어 코멘트 기반 추가된 테스트:
[Copilot on onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc] 및 [xadupre on onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc] 코멘트에서 지적된 것처럼, 새로운 최적화 경로를 검증하기 위한 테스트 케이스가 추가되었습니다. 예를 들어, LayerNorm_Float16_OptimizedRowWithStats와 같은 테스트는 FP16 입력, 16개 이상의 행 크기, 스케일/편향 존재, 그리고 통계 출력(Mean/InvStdDev)을 포함하는 시나리오를 커버하여 새로운 디스패치 로직이 올바르게 작동하는지 확인합니다.
설명:
CPU Provider의 LayerNorm 연산 구현은 입력 데이터 타입, 행 크기, CPU 명령어셋 지원 여부 등을 종합적으로 판단하여 최적의 실행 경로를 선택합니다. 이번 변경으로 FP16 입력이고 CPU가 AVX2 및 F16C를 지원하며 행 크기가 16 이상인 경우, MLAS의 새로운 MlasLayerNormF16 함수를 호출하게 됩니다. 이 함수는 다시 MlasLayerNormKernelF16Avx2와 같은 최적화된 커널로 디스패치됩니다. 이 과정은 기존의 스칼라 기반 FP16 처리 방식보다 훨씬 효율적입니다.
왜 이게 좋은가?
1. 압도적인 성능 향상
PR 설명에 포함된 벤치마크 결과는 이 최적화의 효과를 명확히 보여줍니다.
| Operation | Row size | 1 thread | 8 threads | 16 threads | 32 threads |
|---|---|---|---|---|---|
| LayerNorm | 1024 | 15.66x | 15.39x | 15.85x | 14.89x |
| RMSNorm | 1024 | 12.07x | 11.91x | 11.96x | 10.63x |
| LayerNorm | 4096 | 14.96x | 15.72x | 14.41x | 13.31x |
| RMSNorm | 4096 | 11.71x | 11.42x | 10.84x | 9.61x |
특히 1024 또는 4096과 같이 긴 행 크기에서 LayerNorm은 최대 약 15.85배, RMSNorm은 최대 약 12.07배의 속도 향상을 보였습니다. 이는 기존의 스칼라 기반 FP16 처리 방식이 얼마나 비효율적이었는지를 방증합니다. 긴 시퀀스 길이를 가진 트랜스포머 모델 등에서 이 최적화는 추론 속도를 크게 단축시킬 수 있습니다.
또한, 32 스레드 환경에서 최적화된 커널은 LayerNorm의 경우 초당 54.0-57.6 기가 요소, RMSNorm의 경우 초당 81.2-89.4 기가 요소의 처리량을 달성했습니다. 이는 현대적인 CPU 코어의 잠재력을 최대한 활용하는 수준입니다.
2. 벡터화 및 SIMD 활용의 중요성
이 PR은 현대 CPU의 SIMD(Single Instruction, Multiple Data) 명령어셋, 특히 AVX2, FMA, F16C를 활용하는 것이 고성능 컴퓨팅에서 얼마나 중요한지를 다시 한번 보여줍니다. FP16과 같은 저정밀도 데이터 타입을 사용할 때, 데이터 변환 오버헤드가 성능을 크게 저하시킬 수 있습니다. F16C 명령어는 이러한 변환을 하드웨어 수준에서 매우 빠르게 처리할 수 있게 해주며, AVX2와 FMA는 연산 자체를 병렬적으로 수행하여 처리량을 극대화합니다.
벡터 레지스터를 사용하여 한 번에 여러 데이터를 처리하고, 연산을 FP32로 확장하여 정밀도를 확보한 뒤, 다시 FP16으로 축소하는 전략은 FP16 연산의 이점(메모리 대역폭 절약, 캐시 효율 증대)을 유지하면서도 계산 속도를 크게 향상시키는 효과적인 방법입니다.
3. 동적 디스패치 및 폴백 메커니즘
MlasLayerNormF16 함수가 NormSize < 16 또는 CPU가 F16C를 지원하지 않는 경우 false를 반환하고, platform.cpp에서 F16C 지원 여부를 확인하는 로직은 매우 중요합니다. 이는 모든 환경에서 최적화된 경로를 강제하지 않고, 해당 기능이 지원되지 않거나 이득이 없는 경우에는 자동으로 기존의 안정적인 스칼라 경로로 폴백(fallback)하도록 설계되었기 때문입니다. 이러한 동적 디스패치 및 폴백 메커니즘은 라이브러리의 견고성과 호환성을 보장합니다.
4. 코드 품질 및 테스트의 중요성
리뷰 과정에서 제기된 여러 코멘트들은 코드 품질과 테스트의 중요성을 강조합니다.
- 컴파일러 호환성:
xadupre의 지적처럼, 특정 컴파일러(MSVC)에서 내장 함수 사용 시 발생하는 이슈를 해결하기 위해 플래그를 조정하는 작업은 필수적입니다. - 정확성 보장:
Copilot의 지적처럼,MeanOut이 요청되었을 때 평균 계산을 누락하는 버그를 수정하고, 이를 테스트 케이스로 검증하는 과정은 기능의 정확성을 보장하는 데 중요합니다. - 엔드-투-엔드 테스트:
Copilot과xadupre가 협력하여 추가한 테스트 케이스들은 단순히 MLAS 커널 레벨뿐만 아니라, CPU Provider 레벨에서의 디스패치 및 데이터 변환 로직까지 포함하는 포괄적인 검증을 가능하게 합니다.
이러한 리뷰와 테스트 과정을 통해 최종적으로 안정적이고 성능이 뛰어난 코드가 완성될 수 있었습니다.
결론
Microsoft ONNX Runtime의 이번 PR은 x86 CPU 환경에서 FP16 LayerNorm 및 RMSNorm 연산의 성능을 획기적으로 개선했습니다. AVX2, FMA, F16C 명령어셋을 활용한 벡터화된 연산과 효율적인 데이터 변환은 긴 시퀀스 길이를 가진 딥러닝 모델의 추론 속도를 크게 향상시킬 것입니다. 또한, 동적 디스패치와 철저한 테스트는 이 최적화가 다양한 환경에서 안정적으로 작동하도록 보장합니다. 이는 고성능 라이브러리 개발에서 하드웨어 특성을 최대한 활용하고, 코드 품질을 유지하며, 철저한 검증을 수행하는 것의 중요성을 다시 한번 일깨워주는 좋은 사례입니다.
참고 자료
- https://github.com/microsoft/onnxruntime/blob/main/docs/How_to_use_MLAS.md
- https://en.wikipedia.org/wiki/Advanced_Vector_Extensions#AVX2
- https://en.wikipedia.org/wiki/FMA_instruction_set
- https://en.wikipedia.org/wiki/Half-precision_floating-point_format#FP16_conversion_operations
⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.
관련 포스트
- [onnxruntime] ONNX Runtime: AVX2 및 AVX-VNNI를 위한 2-bit 가중치 CPU 커널 최적화
- [onnxruntime] ONNX Runtime MoE 최적화: QMoE CPU GEMM 성능 대폭 향상
- [onnxruntime] ARM NEON 최적화: LinearAttention 커널 융합으로 3배 성능 향상
- [onnxruntime] ONNX Runtime MLAS, AVX-512 최적화로 MobileClip-S0 모델 추론 속도 향상
- [onnxruntime] ONNX Runtime: Arm64 KleidiAI 기반 FP16 GEMM 및 Convolution 최적화
PR Analysis 의 다른글
- 이전글 [vllm] vLLM GLM5.3 성능 최적화: 메타데이터 연산 속도 1.6~4.8배 향상
- 현재글 : [onnxruntime] ONNX Runtime, x86 CPU에서 FP16 LayerNorm 및 RMSNorm 성능 최적화: AVX2 활용
- 다음글 [flashinfer] NVIDIA Blackwell(SM120)을 위한 초고속 커널 최적화: MiniMax-H3 Fused FC1 + SwiGLU 분석
댓글