What does this PR do ?

Optimizes MoE decode (and small batch prefill) performance by eliminating host synchronizations, enabling end-to-end CUDA graph capture, and using latency-optimized NVLS collectives.

Motivation

The default MoE layer in Megatron-LM is not well-suited for inference. The AlltoAll dispatcher and GroupedGEMM rely on CPU-resident tensors for token-expert assignments, which breaks CUDA graph capture due to host synchronizations. The current workaround — padding to maximum capacity (routing all tokens to all experts) — enables static shapes but wastes significant compute and communication.

This PR introduces an inference-optimized MoE layer that achieves the best of both worlds: compute/communication-optimal routing with full CUDA graph compatibility. These layers are specially designed to optimize the decode phase, where small batch sizes make host synchronization and kernel launch overhead disproportionately expensive.

Optimizations

• InferenceGroupedMLP — Expert computation layer that operates directly on GPU-resident token-expert splits, eliminating the host synchronizations that block CUDA graph capture in the default GroupedGEMM path. Uses FlashInfer cutlass_fused_moe (fused permute + GEMM) for CUDA-graphed iterations and torch._grouped_mm with GPU-resident cumsum offsets for eager mode. Inherits from TEGroupedMLP for checkpoint compatibility..
• InferenceTopKRouter — Stripped-down router that removes training overhead (z-loss, auxiliary losses, token dropping) and is optimized via @torch.compile().
• InferenceCUDAGraphTokenDispatcher — Replaces AlltoAll with AllGather/ReduceScatter for token exchange, keeping all metadata GPU-resident. Supports latency-optimized NVLS collectives on Hopper+ with automatic NCCL fallback.
• Fused 3-tensor NVLS all-gather kernel (multimem_all_gather_fused) for routing_map, probs, and hidden_states — single kernel launch + single barrier.

Other Minor Changes

• InferenceSpecProvider backend for wiring inference-optimized modules into model specs.
• MoELayer dynamically swaps between standard and inference dispatchers based on is_inference_cuda_graphed_iteration.
• Centralized NVLS eligibility checks (are_tensors_nvls_eligible, is_device_nvls_capable) shared across TP and EP communication paths.
• Separate symmetric memory buffer pools for TP and EP (get_global_symmetric_memory_buffer_tp/ep).
• Multi-tensor packing in symmetric memory buffers (maybe_get_tensors) with 16-byte alignment.
• Kill-switch via --inference-disable-triton-nvls-kernels config flag. This makes the system fallback to NCCL.
• Config validation: inference-optimized MoE rejects expert tensor parallelism, capacity-factor routing, and padded routing maps.

How to enable?

these flags -

--transformer-impl inference_optimized \
--moe-router-dtype fp32 # flashinfer only supports fp32 probabilities

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.