NVIDIA / NVIDIA/Megatron-LM

Consider SonicMOE kernel comparison for potentially better IO/activation memory management

Open
#2,709 4 comments 12 reactions 2 assignees Claimed by @sbhavani View on GitHub
community-request enhancement module: moe waiting-on-maintainers
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Is your feature request related to a problem? Please describe.**
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]

SonicMOE has some interesting open source kernels for fine-grained MOE which claim to reduce memory and improves throughput significantly on Hopper. We should see what kernels can be adapted from here for use in TransformerEngine to improve speed and memory usage: https://arxiv.org/abs/2512.14080 Claimed to reduce activation memory by up to 45% improve compute throughput by up to 1.86X on Hopper.

**Describe the solution you'd like**
A clear and concise description of what you want to happen.

Benchmark the kernels and see if there are any easy gains to be had here in TransfromerEngine, specifically with the FusedRouter etc.

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
Add any other context or screenshots about the feature request here.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.