본문으로 건너뛰기

[onnxruntime] ONNX Runtime CUDA NVFP4 QMoE 성능 최적화: 256개 전문가 라우팅 및 GEMV 커널 개선

PR 링크: microsoft/onnxruntime#33545 상태: Merged | 변경: +277 / -36

들어가며

최근 대규모 언어 모델(LLM) 분야에서는 Mixture-of-Experts (MoE) 아키텍처가 큰 주목을 받고 있습니다. 특히, ONNX Runtime은 다양한 하드웨어 가속을 지원하며 LLM 추론 성능 최적화에 힘쓰고 있습니다. 이번 PR은 ONNX Runtime의 CUDA 백엔드에서 NVFP4(NVIDIA FP4 Quantized Mixture-of-Experts) 디코드 경로의 성능을 크게 향상시키는 중요한 개선 사항을 담고 있습니다. 구체적으로는 SM90+ 아키텍처를 타겟으로 256개 전문가를 사용하는 모델에 대한 GEMV(General Matrix Multiply) 연산의 효율성을 높이고, 전문가 라우팅 과정을 더욱 간소화하는 데 중점을 둡니다.

이 PR은 특히 hidden_size = 3072이고 8-토큰 추론을 위한 64개의 확장 행을 사용하는 Ornith-1.5-122B와 같은 대규모 모델에서 성능 병목 현상을 해결하고자 합니다. 기존에는 특정 조건에서 MoeGemvFp4RawKPackedKernel로 폴백되거나, 별도의 커널을 통해 처리되던 라우팅 과정이 새로운 최적화 기법을 통해 통합되고 효율화되었습니다.

코드 분석

이번 PR의 핵심 변경 사항은 크게 두 가지로 나눌 수 있습니다.

  1. MoeGemvFp4DecodeKernel의 기본 동작 범위 확장
  2. 256개 전문가를 위한 단일 CTA 융합 디코드 라우팅 커널 추가

각 변경 사항을 파일별로 자세히 살펴보겠습니다.

1. onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu 및 docs/contrib_ops/cuda/moe_qmoe.md

이 파일들은 NVFP4 GEMV 디코드 커널의 기본 동작 범위를 확장하는 데 관련된 수정 사항을 포함합니다.

변경 전:

-static constexpr int kDecodeDefaultMaxN = 2560;
-static constexpr int kDecodeDefaultMaxK = 2560;
-static constexpr int64_t kDecodeDefaultMaxExpandedRows = 20;

변경 후:

+static constexpr int kDecodeDefaultMaxN = 3072;
+static constexpr int kDecodeDefaultMaxK = 3072;
+static constexpr int64_t kDecodeDefaultMaxExpandedRows = 64;

설명:

kDecodeDefaultMaxN, kDecodeDefaultMaxK, kDecodeDefaultMaxExpandedRows 상수가 증가했습니다. 이전에는 최대 N과 K 차원이 2560, 최대 확장 행(expanded rows)이 20개로 제한되었습니다. 이 PR에서는 이 제한을 각각 3072, 3072, 64개로 늘렸습니다. 이는 hidden_size = 3072와 같이 더 큰 차원을 가진 모델에서 MoeGemvFp4DecodeKernel을 기본적으로 사용할 수 있게 하여, MoeGemvFp4RawKPackedKernel로의 폴백을 줄입니다.

또한, docs/contrib_ops/cuda/moe_qmoe.md 파일에서도 관련 문서가 업데이트되었습니다. ORT_FP4_GEMV_DECODE_KERNEL 옵션에 대한 설명이 변경되어, 기본 활성화 조건이 expanded_num_rows <= 64, N, K <= 3072로 확장되었음을 명시합니다.

-| 	ORT_FP4_GEMV_DECODE_KERNEL | 	1 | 	NVFP4 offline K-packed (`weights_prepacked=1`) decode GEMV mode: `1` (default) enables `MoeGemvFp4DecodeKernel` in the measured SM90 decode envelope (`expanded_num_rows <= 20`, `N, K <= 2560`, and per-projection `|global_scale| * 448 < 3.98`); `0` forces `MoeGemvFp4RawKPackedKernel`; `2` opts into the full kernel envelope (`SM90+`, `expanded_num_rows <= 80`, `K <= 8192`, including the generic scale decoder). Outputs are bit-identical across all modes. Read once per process. See [§9b.2](#9b2-fused-gemv-decode-group-16). |
+|	ORT_FP4_GEMV_DECODE_KERNEL | 	1 | 	NVFP4 offline K-packed (`weights_prepacked=1`) decode GEMV mode: `1` (default) enables `MoeGemvFp4DecodeKernel` in the measured SM90 decode envelope (`expanded_num_rows <= 64`, `N, K <= 3072`, and per-projection `|global_scale| * 448 < 3.98`); `0` forces `MoeGemvFp4RawKPackedKernel`; `2` opts into the full kernel envelope (`SM90+`, `expanded_num_rows <= 80`, `K <= 8192`, including the generic scale decoder). Outputs are bit-identical across all modes. Read once per process. See [§9b.2](#9b2-fused-gemv-decode-group-16). |

2. onnxruntime/contrib_ops/cuda/moe/qmoe_fused_route.cuh 및 관련 파일

이 부분은 256개의 전문가를 지원하는 새로운 융합 라우팅 커널을 추가하고, 기존 라우팅 커널의 동작 범위를 확장하는 내용입니다.

새로운 커널 정의:

+constexpr int kFusedRoute256MaxRows = 8;
+constexpr int kFusedRoute256MaxK = 8;
+constexpr int kFusedRoute256Experts = 256;
+constexpr int kFusedRoute256Threads = 256;
+static_assert(kFusedRoute256MaxRows * kFusedRoute256MaxK <= 64);
+...
+// 256-expert specialization supporting 1..8 rows and k <= 8 (num_rows * k <= 64).
+// Uses 8 warps (256 threads), one warp per row. Lane l owns strided experts j * 32 + l (j = 0..7), ...
+template <typename T, int K>
+__global__ void __launch_bounds__(kFusedRoute256Threads) FusedSoftmaxTopKMaps256Kernel(
+    const T* logits, float* topk_scales, int* topk_indices, int* permuted_row_to_unpermuted_row,
+    int* unpermuted_row_to_permuted_row, int* permuted_expert, int64_t* expert_first_token_offset, int num_rows,
+    int k_runtime, bool normalize_scales) {
+  constexpr int num_experts = kFusedRoute256Experts;
+...
+}

기존 커널 동작 범위 확장:

-| 	ORT_FP4_GEMV_FUSED_ROUTE | 	1 | 	Single-CTA fused softmax + top-k + expert-map decode routing kernel for offline K-packed NVFP4 on SM90+ (`num_rows <= 4`, `num_experts == 512`, `k in {1, 2, 4, 6, 8, 10}`). Bit-identical to `LaunchSoftmaxTopK` + `fusedBuildExpertMapsSortFirstToken`; set to `0` to disable. |
+|	ORT_FP4_GEMV_FUSED_ROUTE | 	1 | 	Single-CTA fused softmax + top-k + expert-map decode routing kernel for offline K-packed NVFP4 on SM90+ (`num_rows <= 8, num_experts == 512, k in {1, 2, 4, 6, 8, 10}` or `num_rows <= 8, num_experts == 256, k in {1, 2, 4, 6, 8}`). Bit-identical to `LaunchSoftmaxTopK` + `fusedBuildExpertMapsSortFirstToken`; set to `0` to disable. |

설명:

FusedSoftmaxTopKMaps256Kernel이라는 새로운 커널이 추가되었습니다. 이 커널은 num_experts = 256인 경우를 위해 특별히 설계되었으며, 최대 8개의 행(num_rows <= 8)과 각 행당 최대 8개의 전문가(k <= 8)를 지원합니다. 이는 num_rows * k <= 64 조건을 만족해야 합니다. 이 커널은 8개의 워프(총 256 스레드)를 사용하여 각 행을 처리하며, 기존 SoftmaxTopKWarpTopKKernel<T, 8, 4>와 동일한 연산 순서로 WarpReduceMax, WarpReduceSum, WarpBitonicTopN 연산을 수행합니다. 이를 통해 전문가 맵(expert maps)을 CTA(Compute Threading Array) 내에서 직접 생성하여, LaunchSoftmaxTopK와 fusedBuildExpertMapsSortFirstToken 두 단계를 거치는 것보다 효율적입니다.

또한, ORT_FP4_GEMV_FUSED_ROUTE 옵션의 설명이 업데이트되어, 이제 256개 전문가와 최대 8개 행, 그리고 k 값의 조합을 지원함을 명시합니다.

리뷰 과정에서 FusedSoftmaxTopKMaps256Kernel의 구현이 Windows CUDA 컴파일러에서 ACCESS_VIOLATION 오류를 발생시키는 문제가 발견되었습니다. 이는 WarpBitonicTopN 구현과 관련이 있었으며, 이후 커밋(b5021b4431)에서 32비트 FusedRouteKey와 __reduce_max_sync (또는 __shfl_xor_sync 폴백)를 사용하는 방식으로 수정되어 문제가 해결되었습니다. 이 수정은 기존 커널들과 동일한 로직을 사용함으로써 비트 일치성(bit-identical)을 유지하면서도 컴파일러 호환성을 확보했습니다.

3. 테스트 및 성능 검증 (onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py)

PR에는 새로운 최적화 기능을 검증하기 위한 테스트 케이스가 추가되었습니다. 특히 N=2048/3072, K=3072/1024, expanded=64 디코드 GEMV 디스패치 및 비트-정확성(bit-exactness) 테스트, 그리고 256개 전문가에 대한 8행 융합 라우팅 및 PDL(Programmatic Dependent Launch) 테스트가 포함되었습니다.

왜 이게 좋은가

이번 PR은 여러 측면에서 중요한 성능 개선을 가져옵니다.

  1. 성능 향상:

    • GEMV 연산: MoeGemvFp4DecodeKernel의 기본 동작 범위를 확장함으로써, 더 많은 경우에 최적화된 커널을 사용할 수 있게 되었습니다. 이는 MoeGemvFp4RawKPackedKernel로의 폴백을 줄여 연산 효율성을 높입니다. H200에서 측정한 결과, 48개 레이어로 구성된 QMoE 모델의 FC1 + FC2 NVFP4 GEMV 연산이 46.2% (11.33ms -> 6.10ms) 빨라졌습니다.
    • MoE 라우팅: 256개 전문가를 위한 새로운 융합 라우팅 커널(FusedSoftmaxTopKMaps256Kernel)은 라우팅 및 전문가 맵 생성 과정을 단일 CTA 내에서 처리하여 오버헤드를 줄입니다. 이 역시 48개 레이어 QMoE 모델에서 64.2% (0.67ms -> 0.24ms)의 성능 향상을 보였습니다.
    • 전체 QMoE 디코드 단계: 결과적으로 48개 레이어 QMoE 디코드 단계에서 총 5.66ms/step의 성능 개선이 이루어졌으며, 이는 약 2배 가까운 속도 향상(12.00ms -> 6.34ms)에 해당합니다.

  2. SM90+ 아키텍처 활용 극대화: 최신 NVIDIA GPU 아키텍처(SM90 이상)의 기능을 최대한 활용하도록 커널을 최적화하여, 최신 하드웨어에서의 성능을 극대화합니다.

  3. 코드 복잡성 감소 및 유지보수성 향상: 기존에는 여러 단계로 나뉘어 처리되던 라우팅 과정을 단일 커널로 통합함으로써 코드의 복잡성을 줄이고 유지보수성을 높였습니다. 또한, MoeGemvFp4DecodeKernel의 적용 범위를 넓혀 폴백 로직을 단순화했습니다.

  4. 일반적인 교훈:

    • 하드웨어 특성 고려: 최신 GPU 아키텍처의 기능을 이해하고 이를 최대한 활용하는 커널 설계가 중요합니다. SM90+의 특성을 고려한 최적화가 좋은 예시입니다.
    • 커널 융합: 여러 단계의 연산을 단일 커널로 융합하는 것은 데이터 이동을 줄이고 GPU 활용률을 높여 성능을 크게 향상시킬 수 있습니다. 라우팅 과정의 융합이 대표적입니다.
    • 동적 범위 확장: 기본 동작 범위를 넓혀 최적화된 경로의 적용률을 높이는 것은 성능 개선의 중요한 전략입니다. GEMV 디코드 커널의 MaxN, MaxK, MaxExpandedRows 확장이 이에 해당합니다.
    • 컴파일러 호환성 및 디버깅: 새로운 커널을 도입할 때는 다양한 플랫폼(특히 Windows)에서의 컴파일러 호환성을 반드시 검증해야 합니다. 리뷰 과정에서 발견된 ACCESS_VIOLATION 문제는 이러한 중요성을 다시 한번 강조합니다.
    • 비트 일치성: 성능 최적화 과정에서 기존 결과와의 비트 일치성을 유지하는 것은 모델의 정확성을 보장하는 데 필수적입니다. 이 PR에서는 이를 엄격하게 검증했습니다.

리뷰 피드백 반영

주요 리뷰 댓글들을 통해 몇 가지 중요한 개선이 이루어졌습니다:

  • 테스트 커버리지 확장: Copilot의 지적대로, FusedSoftmaxTopKMaps256Kernel이 지원하는 다양한 k 값과 행(rows) 조합에 대한 테스트 커버리지를 넓혀, 모든 지원 경로가 철저히 검증되도록 했습니다.
  • Windows 컴파일러 문제 해결: baijumeswani의 지적에 따라, Windows CUDA 12.8 빌드에서 발생하던 ACCESS_VIOLATION 오류를 해결했습니다. 이는 FusedSoftmaxTopKMaps256Kernel의 내부 구현을 수정하여 컴파일러 호환성을 높인 결과입니다. 이 수정으로 인해 관련 빌드 및 문서화 작업도 정상적으로 통과하게 되었습니다.
  • 주석 업데이트: baijumeswani의 지적에 따라, FusedSoftmaxTopKMaps256Kernel의 주석이 실제 구현(bitonic top-k 대신 __reduce_max_sync 기반 arg-max 사용)을 정확하게 반영하도록 수정되었습니다.
  • 구현 및 테스트 수정: tianleiwu가 제시한 문제점들을 해결했습니다. FusedSoftmaxTopKMaps256Kernel의 내부 로직을 수정하여 컴파일러 문제를 해결하고, LaunchFusedRouteMapsT의 커널 디스패치 방식을 개선했으며, softmax_topk_kernel_test.cc에 새로운 테스트 케이스를 추가하여 다양한 시나리오를 검증했습니다.

이러한 리뷰 과정은 코드의 완성도를 높이고 잠재적인 문제를 사전에 발견하여 해결하는 데 크게 기여했습니다.

References

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글