Triton 기반 Fused MoE 커널 공개
Fused MoE dispatch kernel in pure Triton: 89-131% of Megablocks, runs on AMD with zero code changes
·2026.05.27 21:58
Triton으로 구현된 Fused MoE dispatch 커널이 Megablocks 대비 높은 효율과 AMD 호환성을 제공한다.
Triton을 사용하여 CUDA 없이 구현된 Fused MoE dispatch kernel이 공개되었습니다.
이 커널은 추론 배치 사이즈(최대 512 토큰)에서 Stanford의 Megablocks 대비 **89~131%**의 성능을 기록했으며, 코드 수정 없이 AMD MI300X에서도 즉시 구동되는 높은 이식성을 갖췄습니다.
핵심 최적화 기술은 다음과 같습니다:
- Gate+Up projection 융합: SwiGLU 중간 연산 결과가 레지스터를 벗어나지 않도록 설계하여 글로벌 메모리 트래픽을 35% 절감했습니다.
- 커널 실행 최적화: 커널 실행 횟수를 대폭 줄여 오버헤드를 최소화했습니다.
다만, 2048 토큰 이상의 배치 사이즈에서는 Megablocks보다 성능이 낮으며, 라우팅 편향이 심한 64개 이상의 전문가(experts) 환경에서는 성능 저하가 발생하는 한계가 있습니다.
이 요약은 원문 이해를 돕기 위한 큐레이션입니다. 저작권은 원저작자에게 있으며, 정확한 내용과 맥락은 원문을 확인하세요.
요약 오류, 출처 표기 문제, 삭제 요청은 문의 · 건의로 알려주세요.