[ROADMAP][2026 Q2] Megatron Core MoE Roadmap

[Victarry]

Description

The focus for Megatron Core MoE is to provide comprehensive support for latest MoE architectures, advanced parallelism strategies, and performance optimizations for Blackwell. This is a tentative roadmap and subject to change.

• ⚠️ This Roadmap is based on the dev branch; please see the details in its README.
• 🎉 Our technical report for Megatron Core MoE is published in https://arxiv.org/abs/2603.07685
• Previous Roadmaps
  • [ROADMAP][Updated on April 07] Megatron Core MoE Roadmap #1729

---

Megatron MoE Supported Features

Model Support

• ✅ DeepSeek
  • ✅ DeepSeek-V2
  • ✅ DeepSeek-V3, including MTP
  • ✅ DeepSeek-V3.2
• ✅ Qwen
  • ✅ Qwen2-57B-A14B
  • ✅ Qwen3-235B-A22B
  • ✅ Qwen3.5
• ✅ Kimi-K2
• ✅ Mixtral
• 🚀 DeepSeek-V4 (On-going)

Core MoE Functionality

• ✅ Token dropless MoE - Advanced routing without token dropping
• ✅ Top-K Router with flexible K selection
• ✅ Load balancing losses for expert load balancing optimization

Advanced Parallelism

• ✅ Expert Parallel (EP) with 3D parallelism integration
• ✅ Full parallelism combo: EP + DP + TP + PP + SP support
• ✅ Context Parallel (CP) for long sequence MoE training
• ✅ Parallel Folding - Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training
• ✅ Distributed Optimizer for MoE (ZeRO-1 equivalent)
• ✅ (🚀 New!) Megatron FSDP / HSDP with full expert parallel support

Optimizations

• Memory
  • ✅ Memory Efficient token permutation
  • ✅ Pipeline-aware fine-grained activation offloading [Dev] feat(moe): Fine-grained activation offloading #1912
  • ✅ Fine-grained Recomputations (mla, moe, mlp, moe_act, norm)
• Communication
  • ✅ DeepEP support for H100 and B200
  • ✅ (🚀 New!) HybridEP for GB200
  • ✅ 1F1B EP A2A Overlap - Hiding Expert Parallel Communication with 1F1B Pipeline Schedule
  • ✅ DP/PP/TP/EP Communication Overlapping
• Computation
  • ✅ Advanced fusions for Router, Permutation, MLA Rope, FP8 casting, etc.
  • ✅ cuDNN fused Attention and FlashAttn integration
  • ✅ GroupedGEMM and Gradient Accumulation Fusion
  • ✅ Production-ready cudaGraph support for MoE
• Optimizer
  • ✅ Muon and Layer-wise distributed optimizer

Precision Support

• ✅ GroupedGEMM including FP8/MXFP8 support
• ✅ FP8 weights with BF16 optimizer states
• ✅ FP8 training full support
• 🚀 NVFP4 Training

Developer Experience

• ✅ [MoE Model Zoo](https://github.com/yanring/Megatron-MoE-ModelZoo) with pre-training best practices
• ✅ MCore2HF Converter for ecosystem compatibility in megatron-bridge
• ✅ Distributed Checkpointing Support
• ✅ Runtime Upcycling Support for efficient model scaling
• ✅ Layer-wise logging for detailed monitoring

---

(2026 Q2) Megatron MoE Roadmap

Model Supports

DeepSeek-V4

Functionality

• Model Architecture
  • mHC support [dev] feat(mHC): Add basic pytorch implementation of manifold hyper connection(mHC). #2943
  • sqrtsoftplus MoE score function [Dev][MoE] Add a new score function to the router #4193
  • DSA support DeepSeek V3.2 support #2440
  • Hybrid Attention with CSA and HCA [dev] [DeepSeek-v4] Part 1: Hybrid Attention with CSA and HCA #4458
  • Hash MoE, ClampedSwiGLU, mHC contract [dev] [DeepSeek-v4] Part 2: Hash MoE and SwiGLU clamp #4481
  • MTP support with mHC [dev] [DeepSeek-v4] Part 3: MTP support with mHC and new mHC contract #4518
• Muon support
  • Coefficients of Newton-Schulz bump emerging-optimizer package for deepseek v4 coefficients type #4523
  • Distributed Layerwise Optimizer Integrate LayerWiseDistributedOptimizer with DDP buffer infrastructure #4509
  • FP8 primary weight for Muon
  • Precision-aware optimizer
• Packed Sequence
  • DSv4HybridAttention THD support
  • Dispatcher support for THD
• Long Context
  • DSv4HybridAttention CP support
• Convergence
  • HF <-> MCore Checkpoint conversion
  • E2E convergence verification

Performance Optimization

• mHC
  • mHC fusion kernels [dev] mHC kernel fusion  #3828
  • (perf) mHC kernel benchmark and optimization [dev]: faster implementation of mHC fused kernels #4624
  • mHC support with EP Overlapping
• Hybrid Attention
  • Fused Attention Kernel for Dense MQA
  • Fused Attention Kernel for CSA and HCA
  • FP8 Indexer Computation
• ClampedSwiGLU
  • TE unfused path TE#2938
  • Fused kernel in Cudnn-frontend
• batchedGEMM fusion and FP8 support

E2E Training Recipes

• 4K Dense Training
• 16K Dense Training
• 64K Sparse Training
• 1M Sparse Training

Long-Term Topics

• Anticipatory Routing
• Flexible Activation Checkpointing
• MegaMoE

Qwen3.5

Functionality

• Qwen3-Next / Qwen3.5 functionality supported
• Packed sequence support
• HF <-> MCore Checkpoint conversion

Performance Optimization

• Gated Delta Rule (GDN) Optimization
  • GDN kernel fusion on Blackwell (working)
  • GDN FlashQLA integration
  • Pre-gated delta rule kernel fusion (conv1d + memory ops)
• Perf Optimization
  • Full-iteration CUDA graph to reduce CPU overhead of MoE
  • Gated RMSNorm implementation in TE
• Memory Optimization
  • GDN fine-grained activation offloading (P0.5)
  • GDN selective recompute (in_proj, conv1d, gated_delta_rule) (P1)

Long Context Training

• GDN context parallel in A2A/Ulysses mode
• GDN context parallel in AllGather mode — enables Qwen3.5-397B with >64K context (FLA PR, MCore POC PR) (working)
• Dynamic CP support for Multi-modal Training

Multi-Modal Training for Qwen3.5

• OOTB E2E FSDP Example for THD Training Support with Qwen3.5
• Dynamic CP Support
• Heterogeneous Parallelism for Encoder-Decoder

General Optimizations

CUDA Graph

• CUDA Graph Interface Refactor [Dev] Refactor CUDA graph API: decompose cuda_graph_scope into full_iteration impl, inference scope, and per-layer capture modules #4293
• THD + CUDA Graph [Dev] add cuda graph support for thd format training. #4359

Long Context & Context Parallel

• Dynamic CP Support
• Dynamic CP + CUDA Graph

Communication Optimization

• FSDP + A2A Overlap [Dev][feat] Support A2A Overlap for Megatron-FSDP #3796
• DeepEP v2 Token Dispatcher Integration [draft] Add DeepEP v2 flex dispatcher backend #4793
• NCCL-EP Support

Ongoing Long-term Features

• E2E Performance optimization for DeepSeek-V4, Qwen-3.5 and other fine-grained MoEs
• Extreme sparse MoE training exploration
• Migration from GPTModel to HybridModel
• Sync-Free and Full-Iter cudaGraph MoE Training
  • MoE ECHO for dropless MoE load balancing MoE ECHO: Unlocking Sync-Free, Full CUDA-Graph Support for Dropless MoE via Elastic Cloning #2368
• THD and Long Context
  • THD Format E2E Support - End-to-end THD format support [Dev] Add E2E support for THD format #2924
  • Dynamic Context Parallel for Imbalanced Long-Sequence Training
• Megatron FSDP Performance Optimization for MoE Training
• MoE Training with MegaKernels
• Kernel fusions and optimizations for MoE models from TE - MoE training optimization TransformerEngine#2438

(2026 Q1) Highlights

Performance & Kernel Optimizations

• Absorbed MLA - MLA computation optimization for DSA Add absorbed-mla #3193
• Optimization of the standalone permute path Optimization of the standalone permute path deepseek-ai/DeepEP#625
• HybridEP fusion with permutation - Add Permute/Unpermute Fusion Code Path deepseek-ai/DeepEP#588, [Dev] Add permute/unpermute fusion with dispatch/combine in Hybrid-EP #4073
• mHC kernel fusion [dev] mHC kernel fusion  #3828
• GEMM + SwiGLU fused MLP Support GEMM + Swiglu fused MLP #3890

Long Context & Context Parallel

• Hybrid CP Part 2 - Enhanced hybrid data x context parallelism [Dev] feat: Dynamic CP (part 2) #2000
• THD Format E2E Support - End-to-end THD format support [Dev] Add E2E support for THD format #2924

Model & Architecture

• Manifold Hyper Connection (mHC) [dev] feat(mHC): Add basic pytorch implementation of manifold hyper connection(mHC). #2943
• GDN THD Support - Packed sequence support for gated delta net [dev] feat(moe): Support packed sequence for gated delta net (GDN) #2644
• GDN Refinement - Refine gated delta net implementation [dev] perf(moe): Refine gated delta net implementation #3040

Advanced Functionality

• Router replay support for RL training (in progress) feat: add routing replay for Mcore #2693
• Megatron FSDP performance optimization for MoE training

CUDA Graph Enhancements

• Paged Stashing - Dynamic tensor support for dropless MoE with CUDA graph [Dev] Paged Stashing #2690
• CUDA Graph + Offloading - Support CUDA Graph capture with offloading modules [Dev][feat] Support CUDA Graph capture offloading modules #3219
• Optimizer CUDA Graph - Enable CUDA graph for ADAM optimizer Enable CUDA graph for ADAM optimizer #3429

---

Call for Community Contributions

• Model implementations - Additional MoE model variants
• Performance testing - Performance tests across different platforms and workloads
• Documentation and tutorials - Best practices and optimization guides
• Bug fixes

---

This roadmap reflects the collective efforts of NVIDIA and our collaborators.

Credits: MCore MoE Team and @sbhavani

Labels: roadmap, moe, call-for-contribution