[flashinfer] FlashInfer MoE 라우팅 최적화: 더 똑똑해진 전문가 선택으로 성능 향상

2026년 10월 19일수정: 2026년 10월 19일

PR 링크: flashinfer-ai/flashinfer#6555 상태: Merged | 변경: +20 / -16

들어가며

대규모 언어 모델(LLM)의 발전은 모델의 크기와 복잡성을 증가시키는 방향으로 나아가고 있습니다. 특히 Mixture-of-Experts (MoE) 아키텍처는 파라미터 효율성을 높이면서도 뛰어난 성능을 제공하여 주목받고 있습니다. MoE 모델의 핵심은 각 입력 토큰을 어떤 전문가(expert)에게 보낼지 결정하는 라우팅(routing) 메커니즘입니다. 이 라우팅 과정의 효율성은 전체 모델의 성능에 지대한 영향을 미칩니다.

이번 PR (Pull Request)은 FlashInfer 라이브러리에서 MoE 라우팅의 성능을 최적화하는 중요한 변경을 포함하고 있습니다. 기존에는 특정 조건(numExperts >= 160)에서 라우팅 방식이 고정되어 있었지만, 이 PR은 해당 규칙을 재검토하고 실제 워크로드에 더 적합한 동적 라우팅 규칙을 도입하여 성능 병목 현상을 해결하고자 합니다.

코드 분석: 파일별 변경점

`csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cuh`

이 파일은 MoE 라우팅 커스텀 로직을 담당하며, 이번 PR의 핵심 변경 사항이 집중된 곳입니다. 특히 run 함수 내의 라우팅 디스패치 규칙이 수정되었습니다.

1. 라우팅 디스패치 규칙 변경

기존의 디스패치 규칙은 numExperts >= 160이라는 단순한 조건에 기반하여 split-topK 경로를 사용할지 결정했습니다. 하지만 이 규칙은 특정 정책이나 입력 크기 조합에서 최적이 아닐 수 있다는 점이 지적되었습니다.

Before:

-  // Dispatch rule (derived from a policy × (E, K) × BS sweep, see
-  // `bench_routing_sweep.py` + `bench_routing_analyze.py`):
-  //
-  //   useSplit = useSingleCluster && !useSingleBlock && numExperts >= 160
+  // Dispatch rule (derived from a routing-only sweep of every routingCustom policy over
+  // E in {128..1024}, K in {4..22}, BS in [1, 256] on SM100, SM103 and SM107):
+  //
+  //   Few experts (E < 160): split when K >= 8 and BS * K >= 1536, where the fused kernel's
+  //   per-thread K-sized state dominates (up to 2.4x for E=128, K=22).
+  //   Sigmoid+bias policies (DeepSeekV3 nGroup=1, MiniMax2) with E >= 160: always split.
+  //   Other policies with E >= 160: split when BS > 64, or when E >= 512, 10 <= K <= 16 and
+  //   BS >= 32. Below that the fused single-cluster kernel is faster than the scores kernel
+  //   plus the PDL handoff.
   //   plus the bounded high-E/high-K static-block boundary selected above.
-  //
-  // Why this rule (measured on B200 across {Default, Renormalize, 
-  // DeepSeekV3_nGrp1, MiniMax2} policies, BS ∈ [17, 256]):
-  //   - At `numExperts >= 160` split wins on average for every policy:
-  //     1.17× (Renormalize 160/8) up to 3.62× (DeepSeekV3 1024/32).
-  //   - At `numExperts = 128`, split's two-kernel overhead (launch + PDL
-  //     handoff) approximately cancels the single-cluster savings.  Default
-  //     128/8 is ~1.0×; DeepSeekV3 / MiniMax2 128/8 show ~1.11× but
-  //     Renormalize 128/8 is ~1.06× — not worth the complexity.
-  //   - Below the cluster range, block/dynblock remains the default; the bounded high-E/high-K
-  //     static-block boundary above is the measured exception.
-  //
   // Env-var override for benchmarking, applies to *both* the single-cluster
   // split path and the large-BS topK kernel choice:
   //   `FLASHINFER_ROUTING_FORCE_BLOCK_PER_TOKEN`
@@ -1662,8 +1656,18 @@ void run(Data const& data, void* stream) {
         "support block-per-token; the request is ignored.");
   }
 
-  bool useSplitTopKPath = useSingleCluster && !useSingleBlock && policySupportsBlockPerToken &&
-                          (data.mNumExperts >= NumExperts160Experts);
+  bool const splitBeatsFusedCluster = [&] {
+    if (data.mNumExperts < NumExperts160Experts) {
+      return data.mTopK >= 8 && data.mNumTokens * data.mTopK >= 1536;
+    }
+    if (data.mPreprocessType == RoutingPreprocessType::SigmoidBias) {
+      return true;
+    }
+    return data.mNumTokens > 64 || (data.mNumExperts >= 512 && data.mTopK >= 10 &&
+                                    data.mTopK <= 16 && data.mNumTokens >= 32);
+  }();
+  bool useSplitTopKPath =
+      useSingleCluster && !useSingleBlock && policySupportsBlockPerToken && splitBeatsFusedCluster;
   if (forceMode == ForceMode::kOn && useSingleCluster && !useSingleBlock &&
       policySupportsBlockPerToken) {
     useSplitTopKPath = true;

After:

변경된 디스패치 규칙은 다음과 같이 요약할 수 있습니다:

적은 수의 전문가 (E < 160): K >= 8이고 토큰 수 * K >= 1536일 때 split-topK 경로가 유리합니다. 이는 numExperts가 적을 때, 각 스레드에서 처리하는 K개의 상태(state)가 커널 오버헤드를 초과할 때 발생하며, 최대 2.4배의 성능 향상을 보입니다 (예: E=128, K=22).
Sigmoid+bias 정책 (DeepSeekV3 nGroup=1, MiniMax2)과 E >= 160: 이 경우에는 항상 split-topK 경로가 더 빠릅니다.
그 외 정책과 E >= 160: 토큰 수 > 64이거나, E >= 512이고 10 <= K <= 16이며 토큰 수 >= 32일 때 split-topK 경로가 유리합니다. 이 조건들 외에는 기존의 융합된 싱글 클러스터 커널이 점수 커널(scores kernel)과 PDL 핸드오프 비용보다 빠릅니다.

이 새로운 규칙은 기존의 numExperts >= 160이라는 단순한 조건 대신, 전문가 수(E), 선택된 전문가 수(K), 입력 토큰 수(tokens), 그리고 사용되는 전처리 정책(PreprocessType) 등 더 많은 요소를 고려하여 최적의 라우팅 경로를 선택하도록 합니다.

2. `splitBeatsFusedCluster` 람다 함수 도입

변경된 로직은 splitBeatsFusedCluster라는 람다 함수 내에 캡슐화되었습니다. 이 함수는 위에서 설명한 복잡한 조건을 평가하여 split-topK 경로가 기존의 융합된 싱글 클러스터 커널보다 성능상 이점이 있는지 여부를 반환합니다. 이는 코드의 가독성을 높이고 로직을 명확하게 분리하는 역할을 합니다.

왜 이게 좋은가?

이 PR은 다음과 같은 이유로 좋은 최적화 및 개선이라고 할 수 있습니다.

정교해진 라우팅 결정 로직: 기존의 단순한 numExperts >= 160 조건에서 벗어나, 전문가 수, 선택된 전문가 수(K), 토큰 수, 그리고 특정 정책(Sigmoid+bias 등)의 특성을 종합적으로 고려하는 훨씬 정교한 규칙을 도입했습니다. 이는 다양한 워크로드와 모델 설정에 대해 더 나은 성능을 보장합니다.
성능 향상: PR 설명에 따르면, 변경된 규칙은 실제 코드 A/B 테스트에서 다음과 같은 성능 향상을 보였습니다:
- 평균 속도 향상 (Geomean speedup):
  - SM107: 1.106x
  - SM103: 1.080x
  - SM100: 1.078x
- 최대 속도 향상 범위: 0.95x ~ 2.43x (GPU 아키텍처별 상이)
- 특히 DeepSeekV3 nGroup=1 및 MiniMax2 정책에서 1.57x ~ 1.88x의 평균 속도 향상을 기록했습니다.
- 어떤 변경된 경우에도 성능 저하가 발생하지 않았으며, noise band 내에서 유지되었습니다.
코드 가독성 및 유지보수성 향상: 복잡한 라우팅 결정 로직을 splitBeatsFusedCluster 람다 함수로 분리하여 코드의 의도를 명확히 하고, 향후 로직 수정이나 확장을 용이하게 했습니다.
벤치마킹 기반의 과학적 접근: 변경된 규칙은 광범위한 파라미터(E, K, BS)와 다양한 정책에 대한 철저한 라우팅 전용 스윕(sweep) 결과를 기반으로 도출되었습니다. 이는 경험적 추측이 아닌 데이터 기반의 최적화임을 보여줍니다.

https://github.com/flashinfer-ai/flashinfer/blob/main/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cuh

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer SM90 아키텍처 성능 최적화: QkvNormRopeExchangeOverlap의 Shape Gate 개선
현재글 : [flashinfer] FlashInfer MoE 라우팅 최적화: 더 똑똑해진 전문가 선택으로 성능 향상
다음글 없음