[flashinfer] FlashInfer의 Mixture-of-Experts(MoE) 라우팅 성능 최적화 분석

PR 링크: flashinfer-ai/flashinfer#4152 상태: Merged | 변경: +380 / -81

들어가며

최근 대규모 언어 모델(LLM)에서 Mixture-of-Experts(MoE) 구조가 표준으로 자리 잡으면서, 효율적인 Expert 라우팅은 추론 성능의 핵심이 되었습니다. 특히 Expert 수가 많고 TopK 값이 큰 워크로드에서는 기존의 단일 Warp 기반 TopK 정렬 방식이 레지스터 압력과 메모리 대역폭 병목을 유발합니다. 본 PR은 FlashInfer의 TRT-LLM 백엔드에서 이러한 고부하 워크로드를 처리하기 위해 라우팅 정책을 세분화하고, GPU 커널의 병렬 처리 방식을 개선하여 성능을 극대화했습니다.

코드 분석

1. trtllm_fused_moe_routing_custom.cu: 커널 런타임 최적화

기존에는 warp 0가 모든 TopK 정렬을 전담했으나, 새로운 구현에서는 reduceTopKForLane을 도입하여 각 Lane이 자신의 TopK 후보를 유지하도록 변경했습니다.

// Before
BaseType warpTopKScore[KernelParams::MaxNumTopExperts];
int32_t warpTopKExpertIdx[KernelParams::MaxNumTopExperts];
// ... warp 0가 모든 정렬 수행 ...

// After
if constexpr (UseLaneOwnedTopK) {
  ExpertSelect::template applyForLane<...>(...);
} else {
  // 기존 방식 fallback
}

또한, launchBlockKernel에서 numThreadsBlock을 동적으로 계산하여 더 정교한 스케줄링을 지원하도록 수정되었습니다.

2. trtllm_fused_moe_routing_common.cu: 정책 계층 확장

Expert 수가 896개인 경우에 대한 Tier를 추가하여, 기존의 1024/32 Tier로 강제 폴백(fallback)되는 비효율을 제거했습니다.

- bool const launched = routingCustom::launchBlockKernel(customData, numThreadsHist, stream);
+ bool const launched = routingCustom::launchBlockKernel(customData, stream);

왜 이게 좋은가

이번 최적화의 핵심은 '레지스터 압력 완화'와 '병렬성 극대화'입니다.

1. 레지스터 압력 감소: 기존에는 warp 0만 K-sized 레지스터 배열을 들고 있어 병목이 발생했습니다. 이제 각 Lane이 자신의 TopK 후보를 소유함으로써, 블록 전체의 레지스터 사용량이 분산되고 Occupancy가 향상되었습니다.
2. 계층적 병렬화: E=5121024, K=916 구간에서 워커 워프(Worker Warps)들이 128-expert 파티션을 나누어 처리하고, 마지막에 warp 0가 이를 병합하는 방식을 채택했습니다. 이는 대규모 Expert 환경에서 연산 효율을 크게 높입니다.

성능 지표:

• SigmoidBias + ScaledSumNormalize 설정에서 896/16/512 형태의 경우, 기존 대비 약 66.51%의 속도 향상을 기록했습니다.
• 일반적인 고부하 워크로드에서도 10~20% 이상의 유의미한 Latency 감소를 보여줍니다.

교훈

GPU 커널 최적화 시, 특정 Warp에 연산을 집중시키는 것보다 데이터를 Lane 단위로 분산하여 레지스터 압력을 줄이는 전략이 고부하 워크로드에서 훨씬 효과적임을 알 수 있습니다. 또한, 하드웨어의 특성에 맞춘 Tier별 세분화된 커널 디스패치는 범용 커널보다 항상 우수한 성능을 보장합니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#warp-level-primitives
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#launch-bounds

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