[wdykas]

Add MXFP8 batch-invariant inference for Torch/vLLM and FlashInfer, plus TE device-metadata grouped GEMM training. This PR targets main; the selective-precision prerequisite #7300 has already merged.

Summary

• Support MXFP8 batch-invariant inference for Torch and vLLM-selected MXFP8 layers using MCore quantization and PyTorch scaled grouped GEMM. Both squared-ReLU and SwiGLU experts are supported.
• Support FlashInfer routed MXFP8 batch-invariant inference for non-gated squared-ReLU experts, including selective BF16 expert layers.
• Align and zero-fill expert padding before FC1/FC2 quantization; disable fused quantization when the batch-invariant inverse map is needed. Match training-policy SwiGLU probability/rounding order and use 64-bit activation/permutation addressing.
• Use shared RoPE in batch-invariant mode; reject an explicit FlashInfer fused-RoPE request.
• Support TE device-metadata grouped GEMM for training via moe_use_grouped_tensor=True, with 256-row expert padding. Its full workspace is separate for TN/NN/NT layouts and gated to SM100 + cuBLASLt 13.5.1 with TE's GroupedLinear(use_grouped_tensor=...) API. Dense and legacy grouped GEMMs retain restricted workspace. TE op fuser remains unsupported, and the legacy environment switch alone is insufficient.
• Preserve Support selective MXFP8 inference precision #7300's construction-time precision selection and storage-based refit. Native TE inference is outside this PR.

Validation

After merging current main in 860a43c11, the diff contains only the original 17-file change. Sixteen files are identical to the previously reviewed versions; transformer_config.py also preserves the upstream MTP validation changes.

Validation on 860a43c11: six focused configuration tests passed; the determinism coverage gate against main, Black, isort, Ruff, and diff checks passed. The GPU suite below was not repeated for this merge.

Original implementation validation on 691d03eaa:

GB200, PyTorch 2.11.0+cu130, TE 2.20.0.dev0+51317b4, cuBLASLt 13.5.1:

• 984 unit/configuration/FlashInfer/workspace/selective-precision tests passed; 7 optional-dependency tests skipped.
• 24 selected determinism tests passed, including ten actual TE grouped forward/backward/CUDA-graph cases (BF16/MXFP8, changing expert sizes, empty experts, and Nano dimensions). One existing non-strict racy-unpermute negative control XPASS; 17 deselected.
• Six two-GPU EP/NVLS training-vs-inference MoE parity cases passed at atol=0, rtol=0, including TE device-metadata training versus Torch/vLLM-selected MXFP8 inference.
• Four GPT dynamic/batched-forward and co-batch invariance tests passed.
• Black, isort, Ruff, and diff check passed. The determinism coverage gate was rechecked against Support selective MXFP8 inference precision #7300 for the consolidated 17-file diff and passed.

The previous Qwen/Nano NeMo RL zero-logprob smoke results are documented in #7333. Those full-model RL runs were not repeated for this combination; kernel/layer regression results are not a general full-model parity guarantee.

Numerical scope

FlashInfer routed MXFP8 is bitwise invariant across batch composition, but its GEMM is not bitwise identical to the TE training-policy GEMM. The earlier direct two-rank comparison differed in approximately 96% of one-layer BF16 outputs with maximum absolute error 8.5449e-4. Use Torch/vLLM-selected MXFP8 inference for the validated exact training-policy comparisons.

Forward invariance does not promise identical gradients across different minibatches.

[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]

/ok to test af0f048

[wdykas]

/ok to test 1cd8d1f

[wdykas]

/ok to test 76f1d82

[wdykas]

/claude review

[wdykas]

/ok to test 860a43c

[wdykas]

/ok to test 860a43c

[nemo-automation-bot]

🔄 Merge queue validation started!

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