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

[timmoon10]

Description

This PR adds a grouped linear op, which can be used in the grouped MLP block in Mixture-of-Experts models. It also adds an experimental fused operation for a grouped MLP block, using a CuTe DSL kernel that computes an MXFP8 grouped GEMM and SwiGLU.

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

• Add a grouped linear operation
• Add a post-scaled SwiGLU op and add support for interleaving SwiGLU gate and linear units
• Add a fused operation for grouped MLP

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

[timmoon10]

/te-ci pytorch L1

[greptile-apps]

Greptile Overview

Greptile Summary

This PR adds grouped linear operations and experimental fusion for grouped MLP blocks in Mixture-of-Experts models. The implementation includes:

• GroupedLinear: A new operation that applies multiple linear transformations by splitting input along the first dimension, applying separate weights per group, and concatenating results
• ScaledSwiGLU: A post-scaled SwiGLU activation with support for interleaved gate/activation units (32-wide interleaving)
• ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8: An experimental fused operation for SM100+ GPUs using CuTe DSL kernel that combines grouped GEMM, SwiGLU, and post-scaling in a single kernel call

All previously identified issues from earlier review threads have been successfully addressed:

• Fixed undefined group_idx variable usage
• Corrected duplicate condition checks in dimension validation
• Fixed missing f prefixes for f-strings
• Corrected attribute access for per-group weights

The implementation follows established patterns in the codebase, includes comprehensive tests, and properly handles FP8/MXFP8 quantization.

Confidence Score: 5/5

• This PR is safe to merge with no blocking issues found
• All previously identified issues have been resolved, the implementation follows established codebase patterns, includes comprehensive test coverage, and properly handles edge cases
• No files require special attention

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/basic/grouped_linear.py	New GroupedLinear op implementing multiple linear transformations with FP8 quantization support
transformer_engine/pytorch/ops/basic/swiglu.py	New ScaledSwiGLU op with post-scaling and support for gate interleaving
transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py	Experimental fused operation for MXFP8 grouped MLP using CuTe DSL kernel

Sequence Diagram

sequenceDiagram
    participant Input as Input Tensor
    participant FC1 as GroupedLinear (FC1)
    participant SwiGLU as ScaledSwiGLU
    participant FC2 as GroupedLinear (FC2)
    participant Output as Output Tensor
    
    Note over Input,Output: Standard Path (Unfused)
    Input->>FC1: Split by group sizes
    FC1->>FC1: Apply grouped linear transformations
    FC1->>SwiGLU: Interleaved gate/activation output
    SwiGLU->>SwiGLU: Apply SwiGLU + post-scale
    SwiGLU->>FC2: Split activations by group sizes
    FC2->>FC2: Apply grouped linear transformations
    FC2->>Output: Concatenated output
    
    Note over Input,Output: Fused Path (MXFP8 + SM100)
    Input->>FC1: Quantize input (MXFP8)
    FC1->>FC1: CuTe kernel: GEMM + SwiGLU + post-scale
    FC1-->>FC2: Direct MXFP8 intermediate tensors
    FC2->>FC2: Apply grouped GEMM
    FC2->>Output: Final output

Loading

[yaox12]

May be a silly question: are these packing and unpacking code just for verification? Or will they be in the final version?

[timmoon10]

I'm working on getting rid of the concatenations, but the permutes are no-ops. The kernel API expects tensors with non-contiguous dims: https://github.com/NVIDIA/cudnn-frontend/blob/main/python/cudnn/grouped_gemm/grouped_gemm_swiglu/api.py#L240-L245

[skydoorkai]

After swiglu, it usually needs to multiply with permuted_probs. Does this weighted swiglu supported?

[timmoon10]

Yes, the probs are passed into the kernel here: https://github.com/timmoon10/TransformerEngine/blob/46294be478f6551e2cf251283adc7529ddb2964e/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py#L264

[greptile-apps]

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}