[Bug] Wrong gradients with fused grouped MLP under Format.HYBRID (MXFP8)

[Marks101]

Describe the bug

GroupedMLP_CuTeGEMMGLU reinterprets the incoming activation gradient's storage as E4M3 regardless of the recipe's backward format. Under MXFP8BlockScaling(fp8_format=Format.HYBRID) the backward quantizers emit E5M2, so those bytes are reinterpreted rather than converted and every gradient out of the fusion is scaled down by up to 2^-7. The forward pass stays correct, so MoE models train far too slowly instead of failing.

See

TransformerEngine/transformer_engine/pytorch/ops/fused/grouped_mlp.py

Line 1942 in f07a860

fc2_dy_data = grouped_fc2_dy.rowwise_data.view(dtype=data_dtype)

Where

TransformerEngine/transformer_engine/pytorch/ops/fused/grouped_mlp.py

Line 1926 in f07a860

data_dtype = torch.float4_e2m1fn_x2 if use_nvfp4 else torch.float8_e4m3fn

data_dtype is conditional only on NVFP4, never on the FP8 format; e5m2 does not appear anywhere in the file.

Steps/Code to reproduce bug

The existing tests only ever build MXFP8 with Format.E4M3, which is why this isn't caught. Change one line in tests/pytorch/utils.py, in make_recipe's mxfp8 branch:

     if name == "mxfp8":
         return transformer_engine.common.recipe.MXFP8BlockScaling(
-            fp8_format=transformer_engine.common.recipe.Format.E4M3,
+            fp8_format=transformer_engine.common.recipe.Format.HYBRID,

then, on sm100:

NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 pytest \
  "tests/pytorch/test_grouped_mlp.py::TestGroupedMLPFusedOp::test_grouped_mlp" -k mxfp8

Every case that passes with Format.E4M3 fails with Format.HYBRID. The forward assert_close still passes; the failure is in assert_close_grads:

AssertionError: Tensor-likes are not close!
Mismatched elements: 10039 / 393216 (2.6%)
Greatest absolute difference: 0.8675561123053649 at index (436, 214) (up to 0.25 allowed)
Greatest relative difference: 0.9950928443177657 at index (1417, 70) (up to 0.125 allowed)

Expected behavior
Test passes or clear error message.

Environment details
NGC 26.06 PyTorch container, TE 2.17.0, cudnn-frontend 1.24.0

Device details
GB200

Additional context

• The cuDNN grouped GEMMs reject mixed operand dtypes (ValueError: b_dtype (torch.float8_e4m3fn) must match A dtype (torch.float8_e5m2)), so simply passing the true gradient dtype through is not a fix.
• Simplest correct behaviour is to decline the fusion in fuse_grouped_mlp_ops() when get_fp8_torch_dtype(recipe, fprop_tensor=False) isn't E4M3 (falling back to the unfused ops), and to raise rather than reinterpret if such a gradient reaches the kernel path anyway.
• We are aware that Format.E4M3 is the recommended way for MXFP8 ... but still we ran into this bug 🙈