[onnxruntime] ONNX Runtime: Blackwell (SM120+)에서 NVFP4 QMoE를 위한 네이티브 FP4xFP4 Prefill 최적화

PR 링크: microsoft/onnxruntime#29824 상태: Merged | 변경: +406 / -79

들어가며

최근 AI 모델의 크기가 기하급수적으로 커지면서, 모델 경량화와 추론 성능 최적화는 필수적인 과제가 되었습니다. 특히 Mixture-of-Experts (MoE) 모델과 같은 대규모 언어 모델(LLM)에서는 양자화(Quantization) 기법을 활용하여 메모리 사용량을 줄이고 연산 속도를 높이는 것이 중요합니다. NVIDIA의 최신 Blackwell 아키텍처(SM120+)는 FP4와 같은 저정밀도 데이터 타입을 위한 강력한 Tensor Core를 제공하며, 이를 활용하는 것은 성능 향상의 핵심입니다.

이번에 분석할 microsoft/onnxruntime의 PR은 이러한 배경에서 nvfp4 QMoE(Quantized Mixture-of-Experts) 모델의 추론 성능을 획기적으로 개선하고, Blackwell 아키텍처에서 발생하던 빌드 문제를 해결하는 중요한 최적화입니다. 기존에는 nvfp4 양자화 모드가 활성화되어도 FP4 가중치를 FP16/BF16으로 역양자화(dequantize)한 후 A16 dense MoE 러너를 사용했기 때문에 불필요한 오버헤드가 발생했습니다. 이 PR은 Blackwell의 네이티브 FP4xFP4 Tensor Core 연산을 직접 활용하는 새로운 prefill 경로를 도입하여 이 문제를 해결합니다.

또한, SM120 아키텍처에서 LLM object library가 네이티브 sm_120a SASS(Streaming Assembler)를 빌드하지 않아 발생하던 "no kernel image" (CUDA error 209) 오류를 수정하여, Blackwell GPU에서의 ONNX Runtime 안정성을 크게 향상시켰습니다.

코드 분석: 무엇이 왜 좋은 최적화/개선인가

이 PR은 크게 두 가지 핵심 영역에서 변경사항을 가져옵니다: 네이티브 NVFP4 FP4xFP4 prefill 경로 추가와 SM120 빌드 시스템 개선입니다. 각 변경사항을 파일별로 살펴보겠습니다.

1. 네이티브 NVFP4 FP4xFP4 Prefill 경로 추가 (QMoE)

onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc 및 .h

이 파일들은 MoE 양자화 로직의 핵심을 담당하며, 네이티브 FP4xFP4 경로를 활성화하고 라우팅하는 역할을 합니다.

변경 전 (개념):

// moe_quantization.cc (simplified)
// ...
if (quant_params.quant_type == QuantType::kFP4) {
  // Always dequantize FP4 weights to FP16/BF16 and use dense A16 runner
  // construct_dense_a16_runner_();
}
// ...

변경 후:

--- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
+++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
@@ -250,11 +250,19 @@
 
     // Construct the MoE runner
     if (quant_params.quant_type == QuantType::kFP4) {
-      // FP4 (W4A16) always dequantizes to A16 and uses the dense A16 runner.
-      // The fused GEMV / dequant path is used for decode/small-M shapes.
-      // TODO(masahiro): Add native FP4xFP4 path for prefill on SM120+.
-      dense_a16_runner_ = std::make_unique<cutlass::moe::MoERunner<
-          cutlass::arch::Sm80, cutlass::half_t, cutlass::half_t, cutlass::half_t, cutlass::WeightOnlyQuantOp::FP4>>();
+      if (enable_nvfp4_cutlass_gemm_ &&
+          (device_prop.major * 100 + device_prop.minor) >= 120 &&
+          num_rows >= fp4_prefill_min_tokens_ &&
+          (fp4_native_max_tokens_per_expert_ == 0 || num_rows <= fp4_native_max_tokens_per_expert_ * num_experts)) {
+        // Native FP4xFP4 (W4A4) grouped-GEMM prefill path for SM120+.
+        fp4_fp4_runner_ = std::make_unique<cutlass::moe::MoERunner<
+            cutlass::arch::Sm120, cutlass::half_t, cutlass::half_t, cutlass::half_t, cutlass::WeightOnlyQuantOp::FP4,
+            cutlass::ActivationQuantOp::FP4>>();
+      } else {
+        // Fallback to dense A16 (W4A16) runner. Used for decode/small-M shapes or when native path is disabled.
+        dense_a16_runner_ = std::make_unique<cutlass::moe::MoERunner<
+            cutlass::arch::Sm80, cutlass::half_t, cutlass::half_t, cutlass::half_t, cutlass::WeightOnlyQuantOp::FP4>>();
+      }
     }
 // ...

설명:

• 조건부 네이티브 경로 활성화: enable_nvfp4_cutlass_gemm_, GPU 아키텍처(SM120+), 그리고 num_rows (토큰 수)가 fp4_prefill_min_tokens_ 이상인지 여부를 확인하여 네이티브 FP4xFP4 경로를 조건부로 활성화합니다. 이는 prefill 단계에서만 네이티브 경로를 사용하고, decode나 작은 M 값에서는 기존의 dense A16 fallback 경로를 유지하기 위함입니다. fp4_native_max_tokens_per_expert_는 네이티브 경로가 사용될 수 있는 토큰 수의 상한을 설정합니다.
• CUTLASS MoERunner 인스턴스화: 네이티브 경로에서는 cutlass::moe::MoERunner를 cutlass::arch::Sm120 아키텍처와 cutlass::WeightOnlyQuantOp::FP4, cutlass::ActivationQuantOp::FP4를 사용하여 인스턴스화합니다. 이는 FP4 가중치와 FP4 활성화를 모두 사용하여 Blackwell Tensor Core의 FP4xFP4 연산을 직접 활용하겠다는 의미입니다. 기존 fallback 경로는 cutlass::ActivationQuantOp::FP16 (또는 BF16)을 사용했습니다.
• 환경 변수 파싱 개선: getenv 대신 ParseEnvironmentVariableWithDefault를 사용하여 환경 변수 파싱의 견고성을 높였습니다.

왜 좋은가: 이 변경은 nvfp4 QMoE의 prefill 단계에서 불필요한 FP4-to-FP16/BF16 역양자화 오버헤드를 제거합니다. Blackwell GPU의 네이티브 FP4xFP4 Tensor Core 연산을 직접 활용함으로써, 연산 효율성을 극대화하고 prefill 성능을 크게 향상시킵니다. decode 단계에서는 여전히 지연 시간에 민감한 fused GEMV / dequant 경로를 유지하여 전체적인 성능 균형을 맞춥니다.

onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels_fp4_fp4.cu (신규 파일)

이 파일은 FP4xFP4 grouped-GEMM 템플릿의 실제 인스턴스화를 담당합니다.

설명: 이 파일은 CUTLASS 라이브러리를 사용하여 Blackwell 아키텍처에 최적화된 FP4xFP4 grouped-GEMM 커널을 생성합니다. MoE 연산은 여러 전문가(expert)에 대한 GEMM 연산을 그룹화하여 처리하므로, grouped-GEMM은 이러한 구조에 매우 효율적입니다.

왜 좋은가: 특정 데이터 타입(FP4)과 연산(grouped-GEMM)에 최적화된 커널을 직접 인스턴스화함으로써, 하드웨어의 잠재력을 최대한 끌어내어 최고 수준의 성능을 달성할 수 있습니다. 이는 일반적인 GEMM 연산으로는 얻기 어려운 효율성입니다.

onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu

이 파일은 FP4 가중치 및 FP4 활성화 인스턴스화를 지원하도록 업데이트되었습니다.

변경 전 (개념):

// moe_kernels.cu (simplified)
// ... no specific FP4 activation quantization intrinsics

변경 후 (개용):

--- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu
+++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu
@@ -100,6 +100,10 @@
   // ... other code ...
 #if defined(COMPILE_BLACKWELL_SM120_TMA_GROUPED_GEMMS) && defined(USE_FP4_QMOE)
   // Example: FP4 activation quantization using Blackwell-specific intrinsics
+  // The PR description mentions `cvt.e2m1x2` in expandInputRowsKernel,
+  // which is valid only for real sm_120a and cannot be expressed in virtual PTX.
+  // This implies direct use of Blackwell's native FP4 conversion instructions.
+  // ...
 #endif
   // ...
 }

설명: expandInputRowsKernel과 같은 커널에서 cvt.e2m1x2와 같은 Blackwell 고유의 FP4 변환 인트린직(intrinsic)을 사용하여 활성화(activation)를 FP4로 양자화하는 로직이 추가되었습니다. 이는 sm_120a와 같은 실제 아키텍처에서만 유효하며, 가상 PTX로는 표현할 수 없습니다.

왜 좋은가: 활성화까지 FP4로 처리함으로써, W4A4(Weight 4-bit, Activation 4-bit) 연산을 가능하게 하여 메모리 대역폭과 연산량을 더욱 줄일 수 있습니다. 이는 특히 prefill과 같이 대규모 입력이 처리되는 시나리오에서 큰 성능 이점을 제공합니다.

2. 빌드 시스템 / SM120 네이티브 SASS 지원

cmake/onnxruntime_providers_cuda.cmake 및 cmake/CMakeLists.txt

이 파일들은 SM120 아키텍처에 대한 빌드 로직을 수정하여

참고 자료

• https://github.com/NVIDIA/cutlass
• https://www.nvidia.com/en-us/data-center/blackwell-gpu/
• https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html
• https://pytorch.org/docs/stable/generated/torch.compile.html

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