[flashinfer] FlashInfer AlphaMoE W8A8 커널 최적화: (M, K) 기반 동적 커널 디스패치

PR 링크: flashinfer-ai/flashinfer#6512 상태: Merged | 변경: +1821 / -402

들어가며

FlashInfer의 alphamoe_fp8_block_scale_aligned_moe 커널은 NVIDIA B200(sm_100a) 및 GB300(sm_103a) 아키텍처에서 고성능 MoE 연산을 수행합니다. 기존에는 단일 커널 구현으로 다양한 입력 형상(M, K)을 처리했으나, 특정 형상에서 파이프라인 효율이 떨어지는 문제가 있었습니다. 본 PR은 입력 행렬의 크기 (M, K)에 따라 최적화된 세 가지 커널 변형(w8a8_l2, w8a8_i2, w8a8_j2)을 동적으로 디스패치하여 성능을 개선합니다.

코드 분석

csrc/alphamoe_sm100.cu

핵심 변경 사항은 고정된 파이프라인 파라미터를 입력 형상에 따라 분기하는 것입니다. 기존의 단일 커널 구조를 유지하면서, 컴파일 타임에 파이프라인 깊이와 공유 메모리 사용량을 다르게 설정한 세 가지 커널을 생성합니다.

Before (기존 단일 커널):

#define NUM_UP_PIPE_STAGES 3
#define TMEM_NCOLS 72
#define TMEM_DOWN_ACC_OFFSET 48

After (디스패치 구조 도입):

namespace w8a8_i2 {
#define kernel_alpha_moe_w8a8_up_down kernel_alpha_moe_w8a8_up_down_w8a8_i2
#define TMEM_NCOLS 104
#define NUM_UP_PIPE_STAGES 5
#define TMEM_DOWN_ACC_OFFSET 80
// ... 각 형상별로 최적화된 파라미터 정의
}

이 변경은 M과 K의 값에 따라 l2, i2, j2 네임스페이스 내의 커널을 선택적으로 호출하게 하여, 레지스터 압박과 공유 메모리 사용량을 최적화합니다. 특히 M <= 8인 경우 파이프라인 스테이지를 5/3으로 늘려 지연 시간을 숨기는 전략을 취합니다.

왜 이게 좋은가

1. 파이프라인 효율성 극대화: M이 작을 때 파이프라인 스테이지를 늘림으로써(NUM_UP_PIPE_STAGES 5) 메모리 로드와 연산 간의 오버랩을 극대화했습니다.
2. 하드웨어 리소스 최적화: TMEM_NCOLS를 72에서 104로 조정하는 등, 아키텍처별로 가용한 레지스터와 공유 메모리를 더 공격적으로 사용하여 처리량을 높였습니다.
3. 비동기 에필로그(Asynchronous Epilogue): cp.async.bulk.wait_group을 루프 내부로 이동시켜, 출력 블록마다 발생하는 동기화 대기 시간을 줄였습니다.

성능 지표:

• GB300에서 기존 대비 약 1.2배 이상의 성능 향상을 기록했습니다.
• 다양한 (M, K) 조합에서 SOTA 베이스라인(SGLang, TensorRT-LLM)을 상회하거나 동등한 성능을 보입니다.

교훈

고성능 커널 개발에서 'One-size-fits-all' 접근법은 한계가 있습니다. 특히 Hopper/Blackwell 아키텍처와 같이 복잡한 메모리 계층을 가진 환경에서는, 입력 데이터의 형상에 따라 파이프라인 깊이와 공유 메모리 레이아웃을 동적으로 선택하는 'Kernel Specialization'이 필수적입니다. 본 PR은 동일한 소스 코드 패밀리에서 여러 커널을 생성하여 유지보수성과 성능을 모두 잡은 좋은 사례입니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-programming-guide/index.html
• https://github.com/flashinfer-ai/flashinfer

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