[flashinfer] FlashInfer MoE All-to-All 최적화: TRT-LLM의 성능 비결을 파헤치다

PR 링크: flashinfer-ai/flashinfer#3697 상태: Merged | 변경: +1739 / -365

들어가며

최근 대규모 언어 모델(LLM)에서 Mixture-of-Experts(MoE) 아키텍처는 모델의 파라미터 수를 크게 늘리면서도 계산 비용을 효율적으로 유지하는 핵심 기술로 주목받고 있습니다. MoE 모델의 핵심 구성 요소 중 하나는 토큰을 적절한 전문가(Expert)에게 라우팅하고, 각 전문가의 출력을 다시 취합하는 All-to-All(A2A) 통신입니다. 이 과정은 GPU 간의 데이터 전송이 빈번하게 발생하므로, A2A 통신 성능은 MoE 모델 전체의 효율성에 지대한 영향을 미칩니다.

이번에 FlashInfer 저장소에 병합된 PR "Port the TensorRT-LLM one-sided A2A optimizations to Flashinfer"는 NVIDIA TensorRT-LLM에서 검증된 MoE A2A 최적화 기법들을 FlashInfer로 이식하여, MoE 워크로드의 성능과 유연성을 대폭 향상시키는 것을 목표로 합니다. 이 글에서는 해당 PR의 주요 코드 변경사항을 분석하고, 이러한 최적화가 왜 중요한지, 그리고 실제 성능에 어떤 영향을 미치는지 자세히 살펴보겠습니다.

코드 변경사항 분석

이 PR은 크게 네 가지 핵심 영역에서 최적화를 수행했습니다. 각 변경사항을 실제 코드와 함께 살펴보겠습니다.

1. 유연한 Expert-to-Rank 매핑 (Flexible Expert-to-Rank Mapping)

기존 MoE 구현에서는 전문가(expert)들이 GPU 랭크(rank)에 균등하게 분배된다고 가정했습니다. 즉, num_experts가 ep_size (Expert Parallelism size, 즉 랭크 수)로 나누어 떨어지는 경우에만 효율적이었습니다. 하지만 실제 시나리오에서는 num_experts가 ep_size로 나누어 떨어지지 않는 경우가 발생할 수 있으며, 이 경우 부하 불균형이 발생할 수 있습니다. 이 PR은 이러한 문제를 해결하기 위해 compute_target_rank_id 함수를 개선했습니다.

Before:

--- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu
+++ b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu
@@ -193,18 +216,40 @@ __host__ __device__ inline T ceilDiv(T m, T n) {
 // Helper Functions for Expert-to-Rank Mapping
 // ============================================================================
 
-__device__ int compute_target_rank_id(int expert_id, int num_experts_per_rank) {
-  // Compute which rank owns a given expert using contiguous partitioning
-  // Experts are divided evenly across EP ranks:
-  // - Rank 0 gets experts [0, num_experts_per_rank)
-  // - Rank 1 gets experts [num_experts_per_rank, 2*num_experts_per_rank)
-  // - etc.
-  // Example: 32 experts, 4 ranks -> 8 experts per rank
-  // - Rank 0: experts 0-7
-  // - Rank 1: experts 8-15
-  // - Rank 2: experts 16-23
-  // - Rank 3: experts 24-31
-  return expert_id / num_experts_per_rank;
-}

After:

--- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu
+++ b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu
@@ -193,18 +216,40 @@ __host__ __device__ inline T ceilDiv(T m, T n) {
 // Helper Functions for Expert-to-Rank Mapping
 // ============================================================================
 
-__device__ int compute_target_rank_id(int expert_id, int num_experts_per_rank) {
-  // Compute which rank owns a given expert using contiguous partitioning
-  // Experts are divided evenly across EP ranks:
-  // - Rank 0 gets experts [0, num_experts_per_rank)
-  // - Rank 1 gets experts [num_experts_per_rank, 2*num_experts_per_rank)
-  // - etc.
-  // Example: 32 experts, 4 ranks -> 8 experts per rank
-  // - Rank 0: experts 0-7
-  // - Rank 1: experts 8-15
-  // - Rank 2: experts 16-23
-  // - Rank 3: experts 24-31
-  return expert_id / num_experts_per_rank;
+// Compute which rank owns a given expert using contiguous ceil/floor partitioning.
+// Supports non-divisible distribution when num_experts % ep_size != 0:
+//   base      = num_experts / ep_size
+//   remainder = num_experts % ep_size
+//   - Ranks [0, remainder) each own (base + 1) experts.
+//   - Ranks [remainder, ep_size) each own base experts.
+//
+// Example A (uniform): 32 experts, 4 ranks -> base=8, remainder=0
+//   - Rank 0: experts 0-7, Rank 1: 8-15, Rank 2: 16-23, Rank 3: 24-31
+// Example B (non-divisible): 384 experts, 5 ranks -> base=76, remainder=4
+//   - Ranks 0-3: 77 experts each, Rank 4: 76 experts
+//
+// base and remainder are precomputed by the caller once outside the per-token TOP_K loop
+// so the hot path performs at most one integer divide.
+__device__ __forceinline__ int compute_target_rank_id(int expert_id, int base, int remainder) {
+  // Fast path for the uniform (num_experts % ep_size == 0) case: identical to the
+  // pre-ceil/floor implementation, so existing divisible deployments incur no overhead.
+  if (remainder == 0) {
+    return expert_id / base;
+  }
+  int const split = remainder * (base + 1);  // boundary expert id
+  if (expert_id < split) {
+    // Falls inside the (base + 1)-sized prefix block.
+    return expert_id / (base + 1);
+  }
+  // Falls inside the base-sized suffix block.
+  return remainder + (expert_id - split) / base;
+}

새로운 compute_target_rank_id 함수는 num_experts가 ep_size로 나누어 떨어지지 않을 때, remainder개의 랭크에 base + 1개의 전문가를, 나머지 랭크에 base개의 전문가를 할당하는 방식으로 부하를 분산합니다. 이는 MoE 모델의 유연성을 높이고, 다양한 구성에서 더 나은 부하 분산을 가능하게 합니다.

2. Programmatic Launch Dependency (PDL) 지원

Programmatic Launch Dependency (PDL)는 CUDA Hopper 아키텍처(Compute Capability 9.0 이상)에서 도입된 기능으로, GPU 커널 간의 의존성 동기화를 CPU 개입 없이 GPU 자체적으로 관리할 수 있게 해줍니다. 이는 커널 실행 오버헤드를 줄여 전체적인 성능을 향상시킵니다.

Before: (PDL 관련 기능 없음)

After (envUtils.h):

--- a/csrc/nv_internal/tensorrt_llm/common/envUtils.h
+++ b/csrc/nv_internal/tensorrt_llm/common/envUtils.h
@@ -115,4 +119,23 @@ bool getEnvNVFP44Over6ErrUseFastMath();
 // Use 256 instead of 448 for the NVFP4 4over6 E4M3 scaling convention.
 bool getEnvNVFP44Over6E4M3Use256();
 
+template <typename KernelFn, typename... Args>
+inline void launchWithPdlWhenEnabled(char const* name, bool enable_pdl, KernelFn kernelFn,
+                                     dim3 grid, dim3 block, size_t dynamicShmSize,
+                                     cudaStream_t stream, Args&&... args) {
+  cudaLaunchConfig_t kernelConfig;
+  kernelConfig.gridDim = grid;
+  kernelConfig.blockDim = block;
+  kernelConfig.dynamicSmemBytes = dynamicShmSize;
+  kernelConfig.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
+  kernelConfig.attrs = attrs;
+  kernelConfig.numAttrs = 1;
+  cudaError_t e = cudaLaunchKernelEx(&kernelConfig, kernelFn, std::forward<Args>(args)...);
+  FLASHINFER_CHECK(e == cudaSuccess, 

## 참고 자료
- cudaLaunchKernelEx
- cudaLaunchAttributeProgrammaticStreamSerialization
- cudaGridDependencySynchronize
- cudaTriggerProgrammaticLaunchCompletion

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