[flashinfer] FlashInfer, Blackwell GPU를 위한 Gated MoE 커널 최적화로 성능 대폭 향상

PR 링크: flashinfer-ai/flashinfer#4329 상태: Merged | 변경: +7709 / -2639

들어가며

최근 대규모 언어 모델(LLM)의 발전은 Mixture-of-Experts (MoE) 아키텍처의 중요성을 더욱 부각시키고 있습니다. MoE는 모델의 파라미터 수를 크게 늘리면서도 추론 시에는 일부 전문가(expert)만 활성화하여 계산 효율성을 높이는 방식입니다. 하지만 MoE의 효율성은 각 전문가를 얼마나 빠르고 효율적으로 호출하고 계산하느냐에 달려있습니다. 특히, NVIDIA의 최신 Blackwell 아키텍처(SM12x)에서 MoE 연산의 성능을 극대화하기 위한 최적화는 매우 중요합니다.

이번 글에서는 flashinfer-ai/flashinfer 레포지토리의 PR #1137을 분석하여, Blackwell GPU에서 Gated MoE 연산을 위한 커널을 어떻게 최적화했는지 자세히 살펴보겠습니다. 이 PR은 특히 silu, gelu_tanh, swigluoai_uninterleave와 같은 활성화 함수를 사용하는 Gated MoE 연산에서 상당한 성능 향상을 목표로 합니다.

코드 분석

이번 PR의 핵심은 기존의 일반적인 동적 MoE 커널을 Gated MoE 연산에 특화된 최적화된 커널로 분리하고, Blackwell 아키텍처의 특징을 활용하는 것입니다.

1. 커널 분리: generic.py vs gated.py

기존의 동적 MoE 구현은 Gated 연산과 Non-gated 연산을 구분하지 않고 처리했습니다. 하지만 Gated MoE는 output = activation(gate_projection) * up_projection 형태의 연산을 수행하므로, gate와 up 프로젝션 간의 연관성을 활용하여 최적화할 여지가 있습니다. 이 PR은 이를 위해 MoE 동적 커널 구현을 두 파일로 분리했습니다.

• _moe_dynamic/generic.py: 기존의 일반적인 동적 MoE 커널로, Non-gated 활성화 함수(예: relu2)를 처리합니다.
• _moe_dynamic/gated.py: 새로 추가된 최적화된 Gated MoE 커널로, Gated 활성화 함수(silu, gelu_tanh, swigluoai_uninterleave)를 처리합니다.

이 분리를 통해 Gated 연산에 특화된 최적화를 적용할 수 있게 되었습니다.

2. Gated 연산 최적화: Branch Pairing

Gated MoE 연산의 핵심은 gate와 up 프로젝션의 결과를 결합하는 것입니다. 기존 구현에서는 이 두 프로젝션의 계산을 독립적으로 처리하여 스케줄링 오버헤드가 발생했습니다. 새로운 gated.py에서는 이 두 프로젝션을 '쌍(pair)'으로 묶어 N64 FC1 데이터 흐름에서 함께 처리합니다.

PR 설명에 따르면, 이는 다음과 같은 이점을 가집니다:

Pair the gate and up projections in the N64 FC1 dataflow to reduce staging and accumulator lifetimes.

즉, gate와 up 프로젝션의 중간 결과를 저장하고 불러오는(staging) 과정과 누산기(accumulator)의 생명주기를 단축하여 효율성을 높입니다. 이는 Blackwell 아키텍처의 SIMD(Single Instruction, Multiple Data) 특성을 더 잘 활용할 수 있게 합니다.

3. 활성화 함수 지원 및 처리

이 PR은 다음과 같은 활성화 함수를 Gated MoE 커널에서 직접 지원합니다:

• silu (Sigmoid Linear Unit)
• gelu_tanh (Gaussian Error Linear Unit with tanh approximation)
• swigluoai_uninterleave (SwiGLU variant)

특히 silu의 경우, 기존의 튜닝된 explicit inline-PTX reciprocal 경로를 유지하여 성능을 확보했습니다. gelu_tanh와 swigluoai_uninterleave는 표준 활성화 함수 공식을 재사용하면서도 최적화된 Gated 커널 내에서 실행되도록 합니다.

Non-gated 연산(예: FC2)은 기존의 일반적인 실행 경로를 그대로 유지하여 호환성을 보장합니다.

4. NVFP4 양자화 및 데이터 흐름

이 최적화는 NVFP4 (NVIDIA FP4) 양자화를 활용합니다. PR 설명에서 언급된 NVFP4 MoE는 FP4 데이터 타입을 사용하여 모델 가중치를 양자화함으로써 메모리 사용량과 대역폭 요구량을 줄이는 기술입니다. 새로운 Gated 커널은 이러한 NVFP4 양자화된 가중치를 효율적으로 처리하며, 기존의 라우팅, FC2 계산, 스캐터(scatter) 동작을 그대로 유지합니다.

5. API 변경 없음

사용자에게 노출되는 MoEDynamicKernel API는 변경되지 않았습니다. 이는 기존 코드를 사용하는 사용자들에게 영향을 주지 않으면서 내부 구현을 개선할 수 있게 합니다. 또한, 환경 변수 설정 없이 일반적인 활성화 함수 기반 디스패치를 통해 최적화된 경로가 활성화되도록 하여 사용 편의성을 높였습니다.

왜 이게 좋은가?

성능 향상

이 PR의 가장 큰 장점은 명백한 성능 향상입니다. 제공된 벤치마크 결과는 이 최적화가 다양한 모델 설정과 토큰 길이(M)에 걸쳐 상당한 지연 시간 감소를 가져왔음을 보여줍니다.

주요 결과:

• Qwen3.5-35B TP1, M=4096: 1012.692 us -> 586.782 us ( 42.06% 감소)
• Qwen3.5-122B TP2, M=8192: 2517.957 us -> 1441.407 us ( 42.75% 감소)
• 전체 동적 경로 지연 시간 기하 평균 감소: 모델 구성에 따라 8.15% ~ 31.74%

이러한 성능 향상은 Gated MoE 연산의 핵심인 gate와 up 프로젝션의 쌍 처리(branch pairing)를 통해 중간 데이터의 생명주기를 줄이고 스테이징 오버헤드를 최소화한 결과입니다. 또한, Blackwell 아키텍처의 특성을 활용한 최적화된 커널 설계가 주효했습니다.

일반적인 교훈

1. 연산 특화 최적화: 일반적인 커널보다는 특정 연산(여기서는 Gated MoE)의 특성을 깊이 이해하고 이를 반영한 커널을 설계하는 것이 성능 향상에 효과적입니다.
2. 데이터 흐름 분석: gate와 up 프로젝션처럼 연관된 연산들의 데이터 흐름을 분석하여 중간 데이터의 저장 및 접근 패턴을 최적화하는 것이 중요합니다. 이는 레지스터 사용량, 메모리 대역폭, 캐시 효율성에 영향을 미칩니다.
3. 아키텍처 활용: 최신 GPU 아키텍처(Blackwell SM12x)의 새로운 기능이나 특성(예: SIMD 명령어, 특정 데이터 타입 지원)을 적극적으로 활용하는 커널 설계가 필요합니다.
4. API 일관성 유지: 내부 구현을 개선하더라도 사용자에게 노출되는 API는 일관되게 유지하여 하위 호환성을 확보하는 것이 중요합니다.
5. 정확성 검증: 성능 최적화 과정에서 수치적 정확성(NaN, Inf 없음, 허용 오차 내 결과)을 철저히 검증하는 것이 필수적입니다. 이 PR은 BF16 참조 모델과의 비교를 통해 이를 입증했습니다.

리뷰 피드백 반영

리뷰어들의 주요 관심사는 다음과 같았습니다:

• 활성화 함수 기반 디스패치 경계: Generic 커널과 Gated 커널 간의 전환이 활성화 함수에 따라 정확하게 이루어지는지 확인했습니다.
• Gated 동작의 내부 구분: 최상위 Gated 커널 선택과 FC1/FC2 내부의 is_gated 동작 간의 구분을 명확히 했습니다. FC1은 Gated 연산을 수행하고, FC2는 Non-gated 연산을 수행하는 점이 중요합니다.
• SiLU의 explicit inline-PTX reciprocal 경로: SiLU 활성화 함수에서 사용되는 특정 최적화 경로가 올바르게 유지되는지 확인했습니다.
• GELU-tanh 및 SwiGLU-OAI 공식: 최적화된 Gated 커널 내에서 이들 활성화 함수의 공식이 정확하게 구현되었는지 검증했습니다.

이러한 피드백은 코드의 정확성과 안정성을 더욱 높이는 데 기여했습니다.

결론

flashinfer-ai/flashinfer의 이번 PR은 Blackwell GPU 아키텍처에서 Gated MoE 연산의 성능을 획기적으로 개선하는 중요한 작업입니다. 커널을 Gated 연산에 특화시켜 분리하고, gate와 up 프로젝션을 쌍으로 처리하며, NVFP4 양자화를 활용하는 등의 최적화를 통해 상당한 지연 시간 감소를 달성했습니다. 이는 최신 LLM의 효율적인 배포를 위한 중요한 발걸음이며, GPU 커널 최적화의 모범 사례를 보여줍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://github.com/NVIDIA/cutlass/blob/main/docs/ Doxygen/CUDA/classcutlass_1_1gemm_1_1kernel_1_1_sm120_b12x_block_scaled_dense_gemm_kernel.html
• https://github.com/NVIDIA/cutlass/blob/main/docs/ Doxygen/CUDA/namespacecutlass_1_1gemm_1_1kernel_1_1_sm120_b12x.html
• https://github.com/NVIDIA/cutlass/blob/main/docs/ Doxygen/CUDA/namespacecutlass_1_1gemm_1_1kernel_1_1_sm120_b12x_1_1_moe_activation.html

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