Sudhu/core r0.18.0 IFU

[sudhu2k]

What does this PR do ?

Merge in changes for core_r0.18.0 into ROCm

Release Highlights

Parallelism

• A2A overlap for Megatron-FSDP — Bridge Megatron-FSDP's hook-based lifecycle with the EP overlap schedule's direct sub-module execution path ([Main][feat] Support A2A Overlap for Megatron-FSDP NVIDIA/Megatron-LM#3797).
• Packed sequence support for GDN — Packed sequence support for gated delta net on main, following the earlier implementation work ([main] feat(moe): Support packed sequence for gated delta net (GDN) NVIDIA/Megatron-LM#2645, [dev] feat(moe): Support packed sequence for gated delta net (GDN) NVIDIA/Megatron-LM#2644).
• Permute fusion into Hybrid EP — Fuse permute / unpermute with dispatch / combine into a single kernel path for Hybrid EP (add permute fusion into hybrid ep NVIDIA/Megatron-LM#4089).

MoE

• sqrtsoftplus MoE score function — Add a new score function sqrtsoftplus to the router ([MoE] Add a new score function to the router NVIDIA/Megatron-LM#3673).
• Distributed Layerwise Optimizer — Adds a shard-aligned parameter layout for LayerWiseDistributedOptimizer that guarantees no parameter is split across shard boundaries (Integrate LayerWiseDistributedOptimizer with DDP buffer infrastructure NVIDIA/Megatron-LM#4509).
• CUDA Graph Interface Refactor — Decomposes the overloaded cuda_graph_scope field into three dedicated, semantically distinct concepts and cleans up naming throughout (Refactor CUDA graph API: decompose cuda_graph_scope into full_iteration impl, inference scope, and per-layer capture modules NVIDIA/Megatron-LM#4292).

Performance and Memory

• Batch-processing utilities and PP for SFT — Consolidate get_batch utilities and enable pipeline parallel SFT in THD paths (refactor(data): consolidate get_batch and enable PP for SFT THD NVIDIA/Megatron-LM#4103).
• Emerging optimizer refactor — Refactor optimizer infrastructure to generalize beyond Muon and support additional emerging optimizers (Refactor emerging optimizer integration NVIDIA/Megatron-LM#4113, Add more emerging optimizers (#3907) NVIDIA/Megatron-LM#4119).
• Chunked MLP training — Enable chunked MLP during training, extending prior inference-prefill support (Add logic to enable chunked MLP during training NVIDIA/Megatron-LM#3656).
• GroupGEMM + SwiGLU + quantize fused MLP — Transformer Engine fused grouped MLP support, grouped quantized tensor plumbing, checkpoint compatibility, op fuser work, and related MLP infrastructure (Combine GEMM + SwiGLU fused MLP PRs (3890, 4071, 4095, 4219, 4311, 4324) → main NVIDIA/Megatron-LM#4636).
• Full-model CUDA Graph with paged stashing — Decouple oversized compute buffers from properly sized activation buffers for backward-pass storage, continuing the CUDA Graph / dropless MoE memory work (Paged Stashing NVIDIA/Megatron-LM#4247).

Precision

• MXFP8 improvements and fixes — Fine-grained param gather configuration, forced param all-gather eval fixes, numerical fixes when DP overlap is disabled, and follow-through on MXFP8 FP8-param-gather support (fp8-param-gather for mxfp8 NVIDIA/Megatron-LM#2582, M-FSDP: Make fine_grained_param_gather configurable for MXFP8 to enable performance–memory trade-offs NVIDIA/Megatron-LM#4181, [MXFP8/FP4-param-gather] Post processing after forced param AG in eval NVIDIA/Megatron-LM#4562, Fix mxfp8 param gather numerical issue when DP overlap is off NVIDIA/Megatron-LM#4800).
• NVFP4 param gather — Make FP4 param gather work correctly with mixed precision in the NVFP4 recipe (make FP4 param gather work with the mixed precisions in NVFP4 recipe NVIDIA/Megatron-LM#4358).

Inference

• New dispatcher for MoE inference — AllGatherV dispatcher for inference and simplification of the previous dispatcher path (New allgathervdispatcher for inference and simplify old dispatcher.  NVIDIA/Megatron-LM#4258).
• vLLM grouped GEMM backend — Port vLLM grouped GEMM to the inference-optimized MoE backend (Add vLLM grouped gemm backend for MoE inference NVIDIA/Megatron-LM#4566).
• CUDA Graphs for MTP inference — Enable CUDA Graphs for MTP inference (Enable CUDA graphs for MTP inference NVIDIA/Megatron-LM#4260).

Model Architecture and Support

• Flextron — Post-training method for converting a single parent LLM into a nested family of submodels at different parameter budgets (Adding code for Flextron NVIDIA/Megatron-LM#4429).

Multimodal

• MIMO core primitive — Add the core primitive for heterogeneous TP / DP MIMO training, including colocated bridge communication. A step toward the broader MIMO goal of early-fusion multimodal architectures with modular vision, audio, and video encoders, independent nD parallelism per module, and colocated or non-colocated training layouts ([QUESTION] Support for Heterogeneous Parallelism in Multimodal Training NVIDIA/Megatron-LM#1375, Add ColocatedBridgeCommunicator for heterogeneous TP/DP MIMO training (NMFW-17) NVIDIA/Megatron-LM#4368).

Test case results:

https://github.com/ROCm/Megatron-LM/actions/runs/33149814169