순수 Triton으로 구현한 Fused MoE dispatch kernel: Megablocks의 89-131% 성능 구현 및 코드 변경 없이
요약
CUDA 없이 Triton만으로 구현한 Fused MoE dispatch kernel을 소개합니다. Megablocks 대비 89-131%의 성능을 달성했으며, 메모리 트래픽을 35% 절감하고 AMD MI300X에서도 코드 변경 없이 동작합니다.
핵심 포인트
- Triton 기반의 Fused MoE dispatch kernel 개발
- Megablocks 대비 89-131% 수준의 성능 구현
- Gate+Up projection 융합으로 메모리 트래픽 35% 절감
- AMD MI300X 및 NVIDIA A100 호환성 확보
- 대규모 전문가 환경 및 긴 토큰 처리 시 성능 한계 존재
저는 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 성능 측면에 대한 의견을 부탁드립니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 r/OpenAI Codex (search)의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기