[onnxruntime] ONNX Runtime, SM90 GPU를 위한 네이티브 FP8 행렬 곱셈 최적화 도입

PR 링크: microsoft/onnxruntime#32599 상태: Merged | 변경: +698 / -48

들어가며

최근 ONNX Runtime의 PR #32122는 MatMulBlockQuantizedFp8Weight 연산자를 위한 중요한 최적화를 도입했습니다. 이 PR은 특히 NVIDIA의 최신 Hopper 아키텍처(SM90 이상)를 타겟으로 하며, a_scale이 존재하고 block_size=128일 때 네이티브 FP8 행렬 곱셈을 활용하는 새로운 경로를 추가합니다. 이 최적화는 기존의 FP16/BF16 가중치 역양자화 버퍼를 우회하여 상당한 성능 향상을 목표로 합니다.

기존에는 MatMulBlockQuantizedFp8Weight 연산자가 가중치를 FP16 또는 BF16으로 역양자화한 후 cuBLAS를 사용하여 행렬 곱셈을 수행했습니다. 이는 아키텍처에 독립적이지만, 특히 대규모 언어 모델(LLM)과 같이 FP8 양자화가 널리 사용되는 워크로드에서는 추가적인 계산 오버헤드를 발생시켰습니다.

이번 PR은 ORT_FP8_MATMUL_DEEPGEMM=1 환경 변수를 통해 옵트인 방식으로 활성화되는 새로운 경로를 제공합니다. 이 경로는 SM90 GPU의 네이티브 FP8 기능을 활용하여 가중치 역양자화 단계를 건너뛰고 직접 FP8 연산을 수행함으로써, 특히 LLM 추론에서 성능 병목 현상을 완화하는 데 중점을 둡니다.

코드 분석

1. 빌드 시스템 설정 (cmake/ 디렉토리)

이 PR은 DeepGEMM 라이브러리를 ONNX Runtime 빌드에 포함시키기 위해 CMake 빌드 시스템을 수정했습니다. 특히, deep_gemm_matmul_sm90.cu 파일을 LLM 관련 소스 파일 목록에 추가하여 SM90 아키텍처에 특화된 DeepGEMM 커널이 빌드 프로세스에 포함되도록 합니다.

Before:

--- a/cmake/onnxruntime_cuda_source_filters.cmake
+++ b/cmake/onnxruntime_cuda_source_filters.cmake
@@ -166,7 +166,7 @@ function(onnxruntime_extract_llm_sources CU_SRC_LIST)
   set(_llm_fp4_srcs)
   set(_llm_excluded_srcs)
   if(WIN32)
-    list(FILTER _list EXCLUDE REGEX "/moe_gemm/deep_gemm_sm90\.cu$")
+    list(FILTER _list EXCLUDE REGEX "/(moe_gemm/deep_gemm_sm90|deep_gemm_matmul_sm90)\.cu$")
   endif()
   foreach(_src IN LISTS _list)
     if(_src MATCHES "/contrib_ops/cuda/llm/.*\.cu$")
@@ -178,7 +178,7 @@ function(onnxruntime_extract_llm_sources CU_SRC_LIST)
         list(APPEND _llm_excluded_srcs "${_src}")
       # SM90-specific fpA_intB launchers (guarded by #ifndef EXCLUDE_SM_90)
       elseif(_src MATCHES "fpA_intB_gemm_launcher_[0-9]+\.generated\.cu$" OR
-         _src MATCHES "/moe_gemm/deep_gemm_sm90\.cu$")
+         _src MATCHES "/(moe_gemm/deep_gemm_sm90|deep_gemm_matmul_sm90)\.cu$")
         list(APPEND _llm_sm90_srcs "${_src}")
       elseif(onnxruntime_USE_FP4_QMOE AND
              _src MATCHES "/moe_gemm/(moe_gemm_kernels_(bf16|fp16|fp4)_fp4|moe_kernels)\.cu$")

After: deep_gemm_matmul_sm90.cu가 moe_gemm/deep_gemm_sm90.cu와 함께 정규식에 추가되어, 두 파일 모두 LLM 소스 필터링에서 제외되지 않도록 합니다. 이는 해당 파일들이 SM90 관련 커널을 포함하고 있음을 나타냅니다.

이 변경은 onnxruntime_providers_cuda.cmake 및 onnxruntime_providers_cuda_plugin.cmake 파일에서도 유사하게 적용되어, 관련 빌드 구성 요소가 올바르게 포함되도록 합니다.

2. 연산자 문서 업데이트 (docs/ 디렉토리)

MatMulBlockQuantizedFp8Weight 연산자의 동작 방식에 대한 설명이 업데이트되었습니다. 특히, 새로운 옵트인 경로의 존재와 그 동작 방식을 명확히 하기 위해 문서가 수정되었습니다.

Before (ContribOperators.md):

--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -3439,18 +3439,18 @@
 
 ### <a name="com.microsoft.MatMulBlockQuantizedFp8Weight"></a><a name="com.microsoft.matmulblockquantizedfp8weight">**com.microsoft.MatMulBlockQuantizedFp8Weight**</a>
 
-  Weight-only block-scaled FP8 (E4M3) matrix multiplication.
+  Block-scaled FP8 (E4M3) matrix multiplication with optional FP8 activation quantization.
   
-  The weight tensor B is FP8 E4M3 of shape [N, K] with one FP32 scale per `block_size` consecutive
-  K values (`b_scale` of shape [N, ceil(K / block_size)]). The dequantized weight value is
-  `fp8_e4m3(B[n, k]) * b_scale[n, k / block_size]`. The weight is dequantized to the activation
-  type (FP16/BF16) and multiplied with the FP16/BF16 activation A. This path is architecture
-  independent and runs on any CUDA architecture (SM80+).
+  The weight tensor B has shape [N, K] with one FP32 scale per `block_size` consecutive K values
+  (`b_scale` of shape [N, ceil(K / block_size)]). The scaled weight value is
+  `B_scaled[n, k] = fp8_e4m3(B[n, k]) * b_scale[n, k / block_size]`.
   
-  When the optional `a_scale` (a single fp32 scalar) is provided, the activation A is statically
-  quantized to FP8 E4M3 and dequantized back (`a_deq = fp8_e4m3(A / a_scale) * a_scale`) before the
-  matmul, realizing W8A8 activation numerics. When `a_scale` is omitted the activation is kept at
-  full FP16/BF16 precision (weight-only W8A16).
+  When the optional scalar `a_scale` is provided, the activation values used in the multiplication
+  are `A_scaled = fp8_e4m3(A / a_scale) * a_scale` (W8A8). Otherwise, A retains its FP16/BF16
+  precision (weight-only W8A16).
+  
+  The operator multiplies the activation by the transpose of B_scaled and adds the optional bias.
+  The output has shape [..., N] and the same element type as A.
 
 #### Version
 
@@ -3473,7 +3473,7 @@
 <dt><tt>b_scale</tt> : T2</dt>
 <dd>Per-block FP32 weight scales of shape [N, ceil(K / block_size)].</dd>
 <dt><tt>a_scale</tt> (optional) : T2</dt>
-<dd>Optional global fp32 activation scale (scalar). When present, A is statically quantized to FP8 E4M3 with this scale and dequantized back before the matmul (W8A8 numerics); when absent, A stays in full FP16/BF16 precision.</dd>
+<dd>Optional global fp32 activation scale (scalar). When present, A is statically quantized to FP8 E4M3 with this scale (W8A8 numerics); when absent, A retains its FP16/BF16 precision.</dd>
 <dt><tt>bias</tt> (optional) : T</dt>
 <dd>Optional bias of shape [N].</dd>
 </dl>

After (ContribOperators.md): 문서가 업데이트되어 a_scale이 있을 때 W8A8 동작과 없을 때 W8A16 동작을 명확히 구분했습니다. 또한, 새로운 SM90 DeepGEMM 경로에 대한 설명이 docs/contrib_ops/cuda/matmul_block_scaled_fp8.md에 추가되었습니다. 이 문서는 새로운 경로의 활성화 조건, 제약 사항, 동작 방식 및 주의 사항을 상세히 설명합니다.

리뷰 댓글에서 지적된 것처럼, SetDoc 및 a_scale 설명이 기본 경로와 옵트인 경로를 구분하도록 수정되었습니다. 이는 사용자가 새로운 최적화 기능을 올바르게 이해하고 활용하는 데 필수적입니다.

3. 새로운 SM90 DeepGEMM 경로 구현 (matmul_block_scaled_fp8.cc 등)

이 PR의 핵심은 MatMulBlockQuantizedFp8Weight 연산자를 위한 새로운 CUDA 커널 경로입니다. 이 경로는 ORT_FP8_MATMUL_DEEPGEMM=1 플래그가 설정되었을 때 활성화됩니다.

주요 변경 사항은 다음과 같습니다:

• 네이티브 FP8 연산: 가중치를 FP16/BF16으로 역양자화하는 대신, SM90 GPU의 네이티브 FP8 기능을 직접 활용합니다. 이는 sm90_fp8_gemm_1d1d_impl과 같은 DeepGEMM 커널을 호출하여 수행됩니다.
• 가중치 스케일 활용: 기존의 b_scale (K 블록당 FP32 스케일)을 그대로 사용합니다. 새로운 경로는 이 스케일을 FP32 중간 결과에 적용합니다.
• 활성화 양자화: a_scale이 제공되면, 활성화 텐서 A를 FP8 E4M3으로 양자화합니다. 이 과정에서 fp8_e4m3(A / a_scale) * a_scale 연산을 통해 W8A8 연산의 수치적 특성을 모방합니다.
• 출력 처리: FP32 출력은 연산 후 활성화 타입(FP16/BF16)으로 변환되며, 기존의 반올림 순서를 따릅니다. 편향(bias)도 이 변환 후에 추가됩니다.
• 성능 조건: 이 경로는 특정 조건(M<=128, N>=2048, K 및 N의 특정 배수성, 충분한 가중치 크기 등)에서만 활성화됩니다. 이러한 조건은 성능 측정 결과에 기반하며, 해당 범위를 벗어나는 경우에는 기존 경로(GEMV 또는 역양자화 후 cuBLAS)를 사용합니다.

PR 설명에 포함된 벤치마크 결과는 이 새로운 경로가 기존 경로 대비 상당한 성능 향상을 제공함을 보여줍니다.

왜 이게 좋은가?

1. 성능 향상

이 PR의 가장 큰 장점은 SM90 GPU에서 MatMulBlockQuantizedFp8Weight 연산의 성능을 크게 향상시킨다는 점입니다. PR 설명에 제시된 벤치마크 결과는 다음과 같습니다:

M	N	K	Flag off (us)	Flag on (us)	Speedup
64	4096	4096	27.78	17.23	1.61x
128	4096	4096	33.37	19.77	1.69x
64	11008	4096	68.73	31.73	2.17x
128	16384	5120	122.65	63.76	1.92x

이 표는 ORT_FP8_MATMUL_DEEPGEMM=1 플래그를 활성화했을 때, 특히 M과 N 차원이 커질수록 1.6배에서 최대 2.17배까지의 속도 향상을 보여줍니다. 이는 FP8의 네이티브 연산 능력을 최대한 활용하고 불필요한 역양자화 단계를 제거함으로써 달성되었습니다.

2. 메모리 대역폭 및 계산 효율성

기존 경로에서는 가중치를 FP16/BF16으로 역양자화하기 위해 추가적인 메모리 버퍼가 필요했습니다. 새로운 네이티브 FP8 경로는 이 버퍼를 사용하지 않으므로, 메모리 대역폭 사용량을 줄이고 캐시 효율성을 높일 수 있습니다. 또한, FP8 연산은 FP16/BF16 연산보다 일반적으로 더 적은 계산 자원을 요구하므로, GPU의 계산 능력을 더 효율적으로 활용할 수 있습니다.

3. LLM 추론 최적화

대규모 언어 모델(LLM)은 종종 행렬 곱셈 연산에 크게 의존하며, FP8 양자화는 모델 크기를 줄이고 추론 속도를 높이는 데 중요한 기술입니다. 이 PR은 LLM 추론 워크로드에서 병목 현상이 될 수 있는 MatMulBlockQuantizedFp8Weight 연산자를 직접적으로 최적화함으로써, LLM 모델의 전반적인 추론 성능을 향상시키는 데 기여합니다.

4. 옵트인 방식의 장점

새로운 최적화 경로는 기본적으로 비활성화되어 있으며, ORT_FP8_MATMUL_DEEPGEMM=1 플래그를 통해 명시적으로 활성화해야 합니다. 이는 다음과 같은 장점을 제공합니다:

• 안정성: 기본 동작을 변경하지 않으므로 기존 워크로드의 안정성을 보장합니다.
• 유연성: 사용자는 자신의 환경과 모델에 대해 성능 향상이 확실할 때만 이 기능을 활성화할 수 있습니다.
• 정확도 검증: PR 설명에 언급된 바와 같이, 네이티브 FP8 경로는 중간 연산의 반올림 방식이 달라 기존 경로와 비트 단위로 동일한 결과를 보장하지 않을 수 있습니다. 따라서 사용자는 이 기능을 활성화하기 전에 모델의 정확도를 반드시 검증해야 합니다. 옵트인 방식은 이러한 검증 과정을 사용자가 주도적으로 수행할 수 있도록 합니다.

5. 일반적인 교훈

• 하드웨어 특화 최적화: 최신 하드웨어(예: SM90 GPU)의 특정 기능을 활용하는 것은 상당한 성능 향상을 가져올 수 있습니다. ONNX Runtime과 같은 프레임워크는 이러한 하드웨어 기능을 지원하기 위해 지속적으로 발전해야 합니다.
• 양자화의 중요성: FP8과 같은 저정밀도 양자화는 AI 모델의 성능과 효율성을 높이는 데 핵심적인 역할을 합니다. 프레임워크는 이러한 양자화 기술을 효과적으로 지원하고 최적화해야 합니다.
• 옵트인 전략: 새로운 기능이나 성능 최적화를 도입할 때, 옵트인 방식을 채택하는 것은 안정성과 유연성을 동시에 확보하는 좋은 전략입니다. 이는 사용자가 기능의 이점을 누리면서도 잠재적인 위험을 관리할 수 있게 합니다.
• 문서화의 중요성: 복잡한 최적화 기능이나 새로운 경로를 도입할 때는 명확하고 상세한 문서화가 필수적입니다. 사용자가 기능을 올바르게 이해하고, 제약 조건을 파악하며, 잠재적인 영향을 인지하도록 돕는 것은 성공적인 기능 채택에 매우 중요합니다.

References

• Deep Learning Kernel Library (Deep Learning Kernel Library) - 이 PR에서 사용된 DeepGEMM 라이브러리의 GitHub 저장소입니다.
• NVIDIA CUDA Documentation - CUDA 관련 문서입니다.
• ONNX Runtime Contrib Operators Documentation - ONNX Runtime의 기여 연산자들에 대한 공식 문서입니다.
• MatMulBlockQuantizedFp8Weight Operator Documentation - 해당 연산자에 대한 상세 문서 (PR 반영 후 업데이트될 예정).

참고 자료

• https://github.com/microsoft/DeepLearningKernelLibrary
• https://docs.nvidia.com/cuda/
• https://onnxruntime.ai/docs/reference/contrib-ops.html

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