[PyTorch] Decline fused grouped MLP when the backward format is not E4M3

[wilyan09007]

Description

GroupedMLP_CuTeGEMMGLU packs the incoming activation gradient by reinterpreting its storage as E4M3, conditioned only on NVFP4 and never on the FP8 format. Under MXFP8BlockScaling(fp8_format=Format.HYBRID) the backward quantizers emit E5M2, so those bytes are read as the wrong format rather than converted and every gradient out of the fusion is wrong. The forward pass is unaffected.

This declines the fusion when the recipe's backward format is not E4M3, next to the existing RHT check, and raises rather than reinterpreting if such a gradient still reaches the kernel path. NVFP4 pins fp8_format to E4M3, so it is unaffected.

Fixes #3342

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Decline the fusion in fuse_grouped_mlp_ops() when get_fp8_torch_dtype(recipe, fprop_tensor=False) is not E4M3.
• Raise in the fused backward if a non-E4M3 grad output reaches the pack.
• Add test_fusion_declined_for_e5m2_grad_output.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

The PR prevents fused grouped-MLP backward execution when gradients use a format that its packing path would reinterpret incorrectly.

• Declines grouped-MLP fusion when the recipe’s backward FP8 format is not E4M3.
• Adds a runtime safety check before packing non-NVFP4 gradient output.
• Adds a hardware-independent dispatch test covering HYBRID/E5M2 fallback and E4M3 eligibility.

Confidence Score: 5/5

The PR appears safe to merge, with no actionable correctness, security, or test-infrastructure issues identified.

The fusion gate rejects the incompatible HYBRID backward format before rewriting operations, while the runtime check safely catches any non-E4M3 MXFP8 gradient that nevertheless reaches the packing path; supported E4M3 and NVFP4 paths remain enabled.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Adds consistent fusion-time and backward-time safeguards against interpreting E5M2 gradient storage as E4M3.
tests/pytorch/test_grouped_mlp.py	Adds a focused dispatch test proving HYBRID recipes return before fusion scanning while E4M3 recipes proceed.

_{Reviews (1): Last reviewed commit: "[PyTorch] Decline fused grouped MLP when..." | Re-trigger Greptile}