[flashinfer] FlashInfer SM120 MoE GEMM 최적화: 웨이브+잔여물 비용 모델 도입

2026년 8월 12일수정: 2026년 8월 12일

PR 링크: flashinfer-ai/flashinfer#4318 상태: Merged | 변경: +1067 / -122

들어가며

최근 인공지능 모델들은 Mixture-of-Experts (MoE) 아키텍처를 통해 파라미터 수를 폭발적으로 늘리면서도 추론 시에는 일부 전문가(expert)만 활성화하여 효율성을 높이는 방식을 채택하고 있습니다. 이러한 MoE 모델의 핵심 연산 중 하나는 그룹별 행렬 곱셈(Groupwise GEMM)이며, 특히 FP8/MXFP8과 같은 저정밀도 형식을 활용할 때 성능 최적화가 중요합니다.

이번 PR(Pull Request)은 FlashInfer 라이브러리의 SM120 GPU 아키텍처에서 MoE 그룹별 GEMM 연산의 성능을 크게 향상시키는 것을 목표로 합니다. 기존의 단순한 임계값 기반 타일 선택 방식은 특정 상황에서 비효율적인 타일 크기(tile_m)를 선택하여 성능 저하를 야기했습니다. 이 PR은 '웨이브(wave) + 잔여물(residue)' 비용 모델을 도입하고, 기존의 선택 로직을 개선하여 이러한 문제를 해결합니다.

코드 분석

이번 변경사항은 주로 cute_sm120_mxfp8_groupwise/cute_sm120_fp8_op.cu와 cute_sm120_mxfp8_groupwise/cute_sm120_fp8_runner.cu 파일에 집중되어 있습니다. 핵심은 MoE GEMM 연산에서 최적의 tile_m 값을 동적으로 선택하는 로직의 개선입니다.

1. `cute_sm120_mxfp8_groupwise/cute_sm120_fp8_op.cu` 변경사항

이 파일에서는 기존의 CutlassFP8GroupwiseMoeGEMMSM120 함수를 CutlassFP8GroupwiseMoeGEMMSM120Impl로 분리하고, 튜닝된 타일 크기를 전달받을 수 있도록 수정했습니다. 또한, CutlassFP8GroupwiseMoeGEMMSM120Tuned라는 새로운 함수를 추가하여 명시적으로 튜닝된 타일 크기(tactic_tile_m, tactic_tile_n)를 전달받아 사용할 수 있도록 했습니다.

Before:

-void CutlassFP8GroupwiseMoeGEMMSM120(TensorView a, TensorView b, TensorView a_scale,
-                                     TensorView b_scale, TensorView m_indptr, TensorView out,
-                                     std::string scale_major_mode, int64_t scale_granularity_m,
-                                     int64_t scale_granularity_n, int64_t scale_granularity_k,
-                                     int64_t is_gated) {
+static void CutlassFP8GroupwiseMoeGEMMSM120Impl(TensorView a, TensorView b, TensorView a_scale,
+                                                TensorView b_scale, TensorView m_indptr,
+                                                TensorView out, std::string scale_major_mode,
+                                                int64_t scale_granularity_m,
+                                                int64_t scale_granularity_n,
+                                                int64_t scale_granularity_k, int64_t is_gated,
+                                                int64_t tactic_tile_m, int64_t tactic_tile_n) {

After:

-  bool gated = is_gated != 0;
 if (gated) {
     TVM_FFI_ICHECK_EQ(n % 2, 0)
         << "gated (fused SwiGLU) moe requires even b.size(1) (gate+up); got " << n;
@@ -120,10 +132,41 @@ void CutlassFP8GroupwiseMoeGEMMSM120(TensorView a, TensorView b, TensorView a_sc
                                                              float, float>
       runner;
 
-  runner.moe_gemm_fp8_nt_groupwise(
-      out.data_ptr(), static_cast<void const*>(a.data_ptr()),
-      static_cast<void const*>(b.data_ptr()), static_cast<int32_t const*>(m_indptr.data_ptr()),
-      num_experts, total_rows, out_n, k, stream, static_cast<float const*>(a_scale.data_ptr()),
-      static_cast<float const*>(b_scale.data_ptr()), static_cast<int>(scale_granularity_m),
-      static_cast<int>(scale_granularity_n), static_cast<int>(scale_granularity_k), gated);
+  if (tactic_tile_m == -1) {
+    runner.moe_gemm_fp8_nt_groupwise(
+        out.data_ptr(), static_cast<void const*>(a.data_ptr()),
+        static_cast<void const*>(b.data_ptr()), static_cast<int32_t const*>(m_indptr.data_ptr()),
+        num_experts, total_rows, n, k, stream, static_cast<float const*>(a_scale.data_ptr()),
+        static_cast<float const*>(b_scale.data_ptr()), static_cast<int>(scale_granularity_m),
+        static_cast<int>(scale_granularity_n), static_cast<int>(scale_granularity_k), gated);
+  } else {
+    runner.moe_gemm_fp8_nt_groupwise_tuned(
+        out.data_ptr(), static_cast<void const*>(a.data_ptr()),
+        static_cast<void const*>(b.data_ptr()), static_cast<int32_t const*>(m_indptr.data_ptr()),
+        num_experts, total_rows, n, k, stream, static_cast<float const*>(a_scale.data_ptr()),
+        static_cast<float const*>(b_scale.data_ptr()), static_cast<int>(tactic_tile_m),
+        static_cast<int>(tactic_tile_n), static_cast<int>(scale_granularity_m),
+        static_cast<int>(scale_granularity_n), static_cast<int>(scale_granularity_k), gated);
+  }
+}
+
+void CutlassFP8GroupwiseMoeGEMMSM120(TensorView a, TensorView b, TensorView a_scale,
+                                     TensorView b_scale, TensorView m_indptr, TensorView out,
+                                     std::string scale_major_mode, int64_t scale_granularity_m,
+                                     int64_t scale_granularity_n, int64_t scale_granularity_k,
+                                     int64_t is_gated) {
+  CutlassFP8GroupwiseMoeGEMMSM120Impl(a, b, a_scale, b_scale, m_indptr, out, scale_major_mode,
+                                      scale_granularity_m, scale_granularity_n, scale_granularity_k,
+                                      is_gated, -1, -1);
+}
+
+void CutlassFP8GroupwiseMoeGEMMSM120Tuned(TensorView a, TensorView b, TensorView a_scale,
+                                          TensorView b_scale, TensorView m_indptr, TensorView out,
+                                          std::string scale_major_mode, int64_t scale_granularity_m,
+                                          int64_t scale_granularity_n, int64_t scale_granularity_k,
+                                          int64_t is_gated, int64_t tactic_tile_m,
+                                          int64_t tactic_tile_n) {
+  CutlassFP8GroupwiseMoeGEMMSM120Impl(a, b, a_scale, b_scale, m_indptr, out, scale_major_mode,
+                                      scale_granularity_m, scale_granularity_n, scale_granularity_k,
+                                      is_gated, tactic_tile_m, tactic_tile_n);
 }

2. `cute_sm120_mxfp8_groupwise/cute_sm120_fp8_runner.cu` 변경사항

이 파일은 실제 GEMM 연산을 수행하는 러너(runner) 로직을 포함합니다. 가장 중요한 변경은 select_fp8_fused_moe_tile_m와 select_fp8_plain_moe_tile_m 함수의 로직 수정입니다.

기존 로직의 문제점: 기존의 select_fp8_fused_moe_tile_m 함수는 m64_tiles <= num_sms * 8 조건을 통해 num_experts가 클 경우(E >= 64) tile_m=64 옵션을 억제했습니다. 이는 특정 상황에서 최적의 타일 크기 선택을 방해했습니다.

새로운 로직:

select_fp8_fused_moe_tile_m: 이전의 m64_tiles <= num_sms * 8 가드를 제거하여, 더 많은 경우에 tile_m=64 옵션을 고려하도록 했습니다. 이는 '웨이브-어웨어(wave-aware)' 참조 로직과 일치하도록 하여 튜닝된 선택을 가능하게 합니다.
select_fp8_plain_moe_tile_m: 새로운 함수로, select_plain_m64_or_m128 함수를 공유하며 웨이브+잔여물 비용 모델을 기반으로 tile_m을 선택합니다. 이는 기존의 단순 임계값 기반 선택을 대체합니다.

Before (부분 발췌 - select_fp8_fused_moe_tile_m):

-static int select_fp8_fused_moe_tile_m(int total_rows, int shape_n, int num_experts, int num_sms) {
-  int m_per_expert = num_experts > 0 ? (total_rows + num_experts - 1) / num_experts : 0;
+static int select_fp8_fused_moe_tile_m(int total_rows, int shape_n, int shape_k, int num_experts,
+                                       int num_sms) {
+  int max_m_per_expert = sm120_moe_select::balanced_max_rows(total_rows, num_experts);
   auto tile_count = [&](int tile_m, int tile_n) {
-    int64_t num_m = (int64_t(m_per_expert) + tile_m - 1) / tile_m;
-    int64_t num_n = (int64_t(shape_n) + tile_n - 1) / tile_n;
-    return int64_t(num_experts) * num_m * num_n;
+    return sm120_moe_select::balanced_tile_count(total_rows, shape_n, num_experts, tile_m, tile_n);
   };
- 
   int64_t swapab_tiles = tile_count(8, 128);
   int64_t m32_tiles = tile_count(32, 128);
-  int64_t m64_tiles = tile_count(64, 128);
-  int64_t m128_tiles = tile_count(128, 64);
- 
   if (shape_n % 128 == 0 &&
-      (m_per_expert <= 8 || (m32_tiles < num_sms / 2 && swapab_tiles <= num_sms))) {
+      (max_m_per_expert <= 8 || (m32_tiles < num_sms / 2 && swapab_tiles <= num_sms))) {
     return 8;
   }
-  if (shape_n == 64) {
-    return m_per_expert <= 32 ? 32 : 128;
+  if (max_m_per_expert <= 32) {
+    return 32;
   }
-  if (shape_n % 128 != 0) {
-    return 128;
+  if (shape_k <= 2048) {
+    return (max_m_per_expert < 192) ? 64 : 128;
   }
- 
-  int64_t m32_waves = (m32_tiles + num_sms - 1) / num_sms;
-  int64_t m64_waves = (m64_tiles + num_sms - 1) / num_sms;
-  int64_t m128_waves = (m128_tiles + num_sms - 1) / num_sms;
- 
-  if (m64_waves >= m32_waves) {
-    if (m32_waves == 1 && m32_tiles < num_sms / 2) {
-      return 128;
-    }
-    return 32;
+  return sm120_moe_select::select_plain_m64_or_m128(total_rows, shape_n, num_experts, num_sms,
+                                                    /*tile_n_m64=*/128, /*tile_n_m128=*/64);
 }
- if (m128_waves < m64_waves) {
-    return 128;
-  }
- if (m128_waves > m64_waves && m64_tiles <= int64_t(num_sms) * 8) {
-    return 64;
-  }
-  return 128;
+}
+
+static int select_fp8_plain_moe_tile_m(int total_rows, int shape_n, int shape_k, int num_experts,
+                                       int num_sms) {
+  int max_m_per_expert = sm120_moe_select::balanced_max_rows(total_rows, num_experts);
+  auto tile_count = [&](int tile_m, int tile_n) {
+    return sm120_moe_select::balanced_tile_count(total_rows, shape_n, num_experts, tile_m, tile_n);
+  };
+  int64_t swapab_tiles = tile_count(8, 128);
+  int64_t m32_tiles = tile_count(32, 128);
+  if (shape_n % 128 == 0 &&
+      (max_m_per_expert <= 8 || (m32_tiles < num_sms / 2 && swapab_tiles <= num_sms))) {
+    return 8;
+  }
+  if (max_m_per_expert <= 32) {
+    return 32;
+  }
+  if (shape_k <= 2048) {
+    return (max_m_per_expert < 192) ? 64 : 128;
+  }
+  return sm120_moe_select::select_plain_m64_or_m128(total_rows, shape_n, num_experts, num_sms);
 }
 
 static int select_fp8_flat_tile_m(int shape_m, int shape_n, int num_groups, int num_sms) {

3. `cute_sm120_mxfp8_groupwise/sm120_common/moe_tile_selection.h` (추정)

sm120_moe_select::balanced_max_rows, sm120_moe_select::balanced_tile_count, sm120_moe_select::select_plain_m64_or_m128와 같은 헬퍼 함수들이 새로 도입되거나 수정되었습니다. 이 함수들은 각 전문가(expert)에 할당되는 행의 수(max_m_per_expert)를 균형 있게 계산하고, 전체 타일 수를 추정하며, 최종적으로 tile_m을 선택하는 복잡한 로직을 캡슐화합니다. 특히 select_plain_m64_or_m128 함수는 웨이브와 잔여물(residue)의 비용을 모두 고려하여 tile_m=64 또는 tile_m=128을 선택하는 핵심 역할을 합니다.

왜 이게 좋은가?

이 PR은 다음과 같은 이유로 성능 향상에 크게 기여합니다.

정교한 타일 크기 선택: 이전의 단순 임계값 기반 방식은 특정 per-expert-M (각 전문가당 행 수) 범위에서 비효율적인 tile_m을 선택하는 경향이 있었습니다. 예를 들어, per-expert-M이 약 129-207 범위일 때, 기존 방식은 tile_m=128을 선택하여 두 번째 타일의 많은 부분을 낭비했지만, tile_m=64가 더 효율적일 수 있었습니다. 새로운 '웨이브+잔여물' 비용 모델은 이러한 상황을 정확히 파악하여 더 나은 tile_m을 선택합니다.
성능 향상: 실제 벤치마크 결과, 특히 mxfp8 plain의 경우 최대 +38%, fp8 fused의 경우 최대 +17%의 속도 향상을 보여주었습니다. 이는 MoE 모델의 핵심 연산 성능을 직접적으로 개선하여 전체 추론 속도를 높입니다.
- mxfp8 grouped MoE, per-expert-M ≈ 192: 1411 µs -> 1292 µs (+9.2%)
- fp8 fused MoE, per-expert-M ≈ 192: 1507 µs -> 1297 µs (+16.2%)
일반화된 최적화: 이 PR은 특정 시나리오에 국한되지 않고, 다양한 num_sms (스트리밍 멀티프로세서 수)에 적응하며 작동합니다. 또한, per-expert-M이 이미 최적인 경우(예: 256)에는 변경 없이 기존 성능을 유지합니다.
코드 정제 및 유지보수성 향상: select_fp8_fused_moe_tile_m 함수의 불필요한 가드 제거와 cute_sm120_fp8_op.cu에서의 함수 분리 및 튜닝 옵션 추가는 코드의 명확성과 유지보수성을 높입니다.

일반적 교훈:

단순한 임계값은 충분하지 않다: 특히 MoE와 같이 복잡하고 가변적인 워크로드에서는 단순한 임계값 기반의 최적화로는 한계가 있습니다. GPU 아키텍처의 특성(웨이브, 스레드 블록 등)과 연산의 세부 사항(잔여물 처리 등)을 고려한 더 정교한 비용 모델이 필요합니다.
동적 튜닝의 중요성: 하드웨어 및 워크로드 특성에 따라 최적의 파라미터(여기서는 tile_m)가 달라질 수 있습니다. 런타임에 이러한 파라미터를 동적으로 선택하거나 튜닝하는 기능은 성능 극대화에 필수적입니다.
작은 개선의 누적 효과: 각 전문가당 행 수(per-expert-M)의 미묘한 차이로 인해 발생하는 비효율성을 개선하는 것이 전체 성능에 상당한 영향을 미칠 수 있습니다. 이는 GPU 커널 최적화에서 세밀한 부분까지 신경 써야 함을 보여줍니다.

리뷰 피드백 반영

리뷰어(CarstyYou)는 tile_m 선택 로직의 정확성을 높이기 위해 다음과 같은 사항을 제안하고 수정되었습니다:

Plain FP8/MXFP8 선택기 개선: 각 전문가의 균형 잡힌 행 수를 정확히 계산하도록 수정되었으며, 나머지 행까지 포함하도록 개선되었습니다.
FP8 Fused Fallback 로직: 논리적인 out_n 값을 전달하고, TileN을 올바르게 모델링하도록 수정되었습니다. (M32/M64/SWAPAB는 TileN=128, M128은 TileN=64 사용)

이러한 피드백은 코드의 정확성을 높이고 다양한 시나리오에서 일관된 성능을 보장하는 데 기여했습니다.

References

FlashInfer GitHub Repository
SM120 Common Utilities (추정 - 실제 파일 경로 및 내용 다를 수 있음)
[Cutlass GEMM Kernels](https://github.com/ NVIDIA/cutlass) (FlashInfer가 기반으로 하는 라이브러리)

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

PR Analysis 의 다른글

이전글 [sglang] [Diffusion] MiniMax-H3 비디오 인입 최적화: 87.5% 메모리 절감과 Zero-Copy의 마법
현재글 : [flashinfer] FlashInfer SM120 MoE GEMM 최적화: 웨이브+잔여물 비용 모델 도입
다음글 [sglang] SGLang의 FLUX.2 추론 최적화: Eager 모드에서의 커널 융합 기법