[CharlesXu-HQ]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Fuses batch-invariant MCore MoE token permutation with MXFP8 quantization while preserving the inverse-map and zero-padding contracts required by deterministic unpermute.

The fused Triton kernel can now optionally:

• emit the (token, local expert) -> permuted row inverse map;
• initialize MXFP8 alignment-padding rows to zero; and
• retain the existing four-value API when neither option is requested.

mcore_fused_moe uses this path for batch-invariant MXFP8 squared-ReLU and SwiGLU. Activation behavior is unchanged; only input permutation and quantization are fused.

Issue tracking

Linked issue: Fixes #7659

Performance

RTX 5090, PyTorch 2.11.0+cu128, Triton 3.6.0. Each measurement covers batch-invariant preprocessing: permutation, zero padding, inverse-map generation, and MXFP8 quantization.

tokens	hidden	top-k	local experts	separate	fused	speedup
128	2688	8	128	0.121 ms	0.045 ms	2.70x
512	4096	8	64	0.121 ms	0.047 ms	2.56x
2048	4096	8	16	0.260 ms	0.087 ms	2.98x

Validation

• TestPermuteAndQuantizeMxfp8 plus test_batch_invariant_mxfp8_permute_replays: 24 passed, 1 environment-incompatible test deselected on SM120.
• Fused vs. separate preprocessing is byte-exact for logical FP8 data, E8M0 scales, probabilities, inverse maps, and alignment padding.
• Squared-ReLU and SwiGLU dispatch tests exercise the real fused preprocessing path; grouped GEMM is stubbed because PyTorch's MXFP8 scaled_grouped_mm does not support SM120.
• tools/check_kernel_determinism_coverage.py --base-ref upstream/main passes.
• Black, isort, Ruff, py_compile, and git diff --check pass for all changed files.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added or updated the bit-exact determinism test; the existing inference_moe_kernels manifest entry covers this kernel file
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter checks on my PR

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[wdykas]

thanks for contributing! left a small comment but looks good overall.

[wdykas]

/ok to test 213b1b8

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/37327446671