저는 MoE (Mixture-of-Experts) 추론 작업을 진행해 왔으며, CUDA 없이 완전히 Triton으로 작성된 fused dispatch kernel을 개발했습니다.

추론 배치 크기(최대 512 토큰)에서 Megablocks(Stanford의 CUDA 최적화 MoE 라이브러리)의 89-131% 성능에 도달하며, 동일한 커널이 코드 변경 없이 AMD MI300X에서도 실행됩니다. A100에서 Mixtral-8x7B를 대상으로 테스트했습니다.

가장 큰 성과는 gate+up projection을 fusing(융합)하여 SwiGLU 중간값이 레지스터(registers)를 벗어나지 않도록 한 것이며, 이를 통해 글로벌 메모리 트래픽(global memory traffic)을 35% 절감했습니다. 커널 실행(kernel launches) 횟수를 줄인 것(24회 이상에서 5회로)도 도움이 되었지만, 그 비중은 상대적으로 적었습니다.

솔직한 한계점은 다음과 같습니다: 2048 토큰 이상에서는 Megablocks에 뒤처지며, 심한 라우팅 왜곡(routing skew)이 발생하는 64개 이상의 전문가(experts) 환경에서는 여전히 성능이 저조합니다. 따라서 DeepSeek-V3 규모의 전문가 수에는 아직 미치지 못합니다.

코드: https://github.com/bassrehab/triton-kernels

벤치마크가 포함된 기술 문서: https://subhadipmitra.com/blog/2026/fused-moe-dispatch-triton/

논문: https://arxiv.org/abs/2605.23911

피드백을 환영합니다. 특히 아직 최적화되지 않은 AMD 성능 측면에 대한 의견을 부탁드립니다.