Use fused pad+MCT when available

[alextmagro]

fuses the padding kernels and cast transpose kernels that are used for FP8 MoE operations. Fuses both in forward and backward pass.

[sudhu2k]

Otherwise looks good to me.

[sudhu2k]

Use the logger object from line 72 to do
Use logger.warning(
"message"
) instead

[Copilot]

Pull request overview

This PR adds an ROCm/TransformerEngine-gated fast path for FP8/FP4 MoE that fuses padding with TE’s cast+transpose (MCT) in grouped GEMMs, aiming to avoid BF16 intermediate padding and optionally unpad FC2 outputs.

Changes:

• Detect TE support for fused padding (via Fp8Padding.compute_padded_splits) and use it to compute padded m_splits without materializing a BF16-padded activation.
• Plumb actual_m_splits (and unpad_output for FC2) through TEGroupedMLP into the TE GroupedLinear wrapper when supported by the installed TE version.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
megatron/core/transformer/moe/experts.py	Adds ROCm fused padding path and passes actual_m_splits / unpad_output into grouped linears.
megatron/core/extensions/transformer_engine.py	Extends TE GroupedLinear wrapper signature and conditionally forwards fused-padding kwargs to TE.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[sudhu2k]

We don't have a fp8 triton groupedgemm kernel for now and we would fall back to hipblasLT Multistream GEMM. But I guess it would be future proof to do this and the better way would be to pad on the gpu tensors directly instead of creating a new tensor from a list.

[alextmagro]

I think maybe this should be postponed until we have intentions to do a grouped gemm in triton.