본문으로 건너뛰기

[flashinfer] FlashInfer의 TRTLLM-Gen MoE 라우팅 최적화: 레지스터 압박 해소와 성능 극대화

PR 링크: flashinfer-ai/flashinfer#3751 상태: Merged | 변경: +292 / -51

들어가며

최신 대규모 언어 모델(LLM)에서 Mixture-of-Experts(MoE) 아키텍처는 필수적인 요소가 되었습니다. 하지만 MoE의 라우팅 단계는 많은 전문가(Expert)와 토큰을 처리할 때 레지스터 압박(Register Pressure)으로 인한 성능 저하가 빈번하게 발생합니다. 특히 고정된 스레드 블록 크기를 사용할 경우, 작은 워크로드에서는 리소스 낭비가, 큰 워크로드에서는 과도한 스필링(Spilling)이 발생합니다. 본 PR은 flashinfer-ai/flashinfer의 TRTLLM-Gen MoE 라우팅 커널을 개선하여, 워크로드 크기에 따라 스레드 블록 크기를 동적으로 선택함으로써 성능을 최적화했습니다.

코드 분석

1. 동적 스레드 블록 크기 도입 (trtllm_fused_moe_routing_custom.cu)

기존에는 모든 상황에서 고정된 스레드 블록 크기를 사용했으나, 이번 변경으로 256, 512, 1024 스레드 블록 크기를 지원하는 템플릿 커널이 추가되었습니다.

Before:

// 기존에는 고정된 NumThreads를 사용
__global__ void __cluster_dims__(NumBlocksPerCluster, 1, 1) __launch_bounds__(NumThreads)
    routingIndicesClusterKernel(KernelParams params) { ... }

After:

// 템플릿을 통해 ClusterBlockDim을 동적으로 설정
template <typename KernelParams, int ClusterBlockDim = NumThreads>
__global__ void __cluster_dims__(NumBlocksPerCluster, 1, 1) __launch_bounds__(ClusterBlockDim)
    routingIndicesClusterKernel(KernelParams params) { ... }

이 변경을 통해 커널은 mNumTokens에 따라 적절한 블록 크기를 선택하여 실행됩니다. 예를 들어, 토큰 수가 적을 때는 더 작은 블록을 사용하여 레지스터 사용량을 최적화하고, 토큰 수가 많을 때는 더 큰 블록을 사용하여 병렬성을 극대화합니다.

2. 라우팅 정책의 유연성 확보

ClusterPolicyTraits와 FilterClusterTiers 구조체를 도입하여, 특정 전문가 수와 블록 크기 조합이 유효한지 컴파일 타임에 검증합니다. 이는 불필요한 커널 생성을 방지하고, 하드웨어 아키텍처에 맞는 최적의 경로를 선택하게 합니다.

왜 이게 좋은가

이번 최적화의 핵심은 '레지스터 압박 완화'입니다. GPU 커널에서 레지스터 사용량이 많아지면 스레드당 가용 레지스터가 줄어들어 스필링이 발생하고, 이는 메모리 대역폭을 낭비하게 됩니다. 블록 크기를 줄임으로써 각 스레드 블록이 사용하는 레지스터를 최적화하고, 결과적으로 더 높은 점유율(Occupancy)을 확보할 수 있습니다.

성능 수치

제공된 벤치마크 결과에 따르면, 특히 1024 전문가 설정에서 softmax_sum 연산의 경우 최대 12배 이상의 속도 향상을 보였습니다. 이는 고정된 커널이 처리하기 힘들었던 복잡한 라우팅 패턴을 워크로드 크기에 맞게 세분화한 결과입니다.

교훈

  1. 워크로드 기반의 휴리스틱: 모든 상황에 맞는 단일 커널은 없습니다. 토큰 수나 전문가 수에 따라 커널 실행 설정을 분기하는 것이 필수적입니다.
  2. 컴파일 타임 최적화: C++ 템플릿 메타프로그래밍을 활용하여 유효한 커널 경로만 생성함으로써 런타임 오버헤드를 최소화할 수 있습니다.
  3. 레지스터 관리: 고성능 CUDA 커널 작성 시, __launch_bounds__와 블록 크기 조절을 통해 레지스터 스필링을 제어하는 것이 성능의 핵심입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글