[PyTorch] Add grouped linear op and experimental fusion for grouped MLP

tests/pytorch/test_fusible_ops.py

@@ -3380,6 +3380,27 @@ def test_grouped_mlp(
             y_test = module(x_test, split_sizes, probs_test, split_sizes)
         y_test.backward(dy_test)
 
+        # Check for expected fusions
+        if (
+            te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported()
+            and quantization == "mxfp8"
+            and dtype == torch.bfloat16
+            and not bias
+            and glu_interleave_size == 32
+        ):
+            forward_ops = module._module_groups[0]._forward_ops
+            backward_ops = module._module_groups[0]._backward_ops
+            assert len(forward_ops) == 1
+            assert isinstance(
+                forward_ops[0][0],
+                te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8,
+            )
+            assert len(backward_ops) == 1
+            assert isinstance(
+                backward_ops[0][0],
+                te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8,
+            )
+
         # Loose tols for sanity checking
         tols = {"rtol": 0.125, "atol": 0.25}
         if quantization == "nvfp4":

transformer_engine/pytorch/ops/fused/__init__.py

@@ -28,3 +28,12 @@
 register_backward_fusion(BackwardLinearScale.fuse_backward_ops)
 register_backward_fusion(BackwardActivationBias.fuse_backward_ops)
 register_backward_fusion(BackwardAddRMSNorm.fuse_backward_ops)
+
+# Import experimental fusions
+# Note: Registration logic is non-trivial, so submodule handles it internally.
+from .forward_grouped_mlp import (  # pylint: disable=wrong-import-position
+    ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8,
+)
+from .backward_grouped_mlp import (  # pylint: disable=wrong-import-position
+    BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8,
+)