Activation + GroupedLinear Fusion for MOE and other MOE optimizations

[vthumbe1503]

Description

Better Kernels for Scaled Activations

• Scaled Srelu, Swiglu and Clamped Swiglu Pytorch Ops are replaced with their corresponding tex APIs introduced in this PR. tex APIs under the hood uses fused kernels introduced in this PR
• Implementaion of the tex APIs are done in activation.cpp

Activation + GroupQuantize Fusion

• Similar tex APIs for scaled activations, there are tex APIs created for grouped_scaled_* activations. Currently they are implemented by callng nvte APIs for scaled activations followed by the group quantize.
• Eventual goal is to replace them with the fused kernels. But at the moment they will serve to reduce CPU overheads where in we will have to just call one single grouped_scaled_* API instead of calling the 2 pybinded tex APIs for scaled_* activation and then group_quantize.
• Fusion are implemented in ops via pairwise fusions between ScaledActivation and GroupedLinear ops. cc: @timmoon10

Precomputed Tensor Offsets Optimization

• Tensor Offsets were already precomputed in forward pass via tex API splits_to_offset_multi. The same is done now for backward pass as well and this saves use two wasted kernel calls for calculation of offsets in the backward pass in grouped linear and 4 wasted kernel calls in grouped mlp.

BF16 Grouped MLP Performance for num_groups=8 and swiglu activation

Metric	GPU-Bound Case (hidden_dim = 2048)	CPU-Bound Case (hidden_dim = 1024)
Tokens	65,536	2,048
Before PR	5.722 ms / iter	1.385 ms / iter
After PR	5.329 ms / iter	1.047 ms / iter
Delta (%)	🟢 -6.87% (Faster)	🟢 -24.40% (Faster)

MXFP8 Grouped MLP performance for num_groups = 8 and swiglu activation

With CuteDSL fusions

Metric	GPU-Bound Case (hidden_dim = 2048)	CPU-Bound Case (hidden_dim = 128)
Tokens	65,536	2,048
Before PR	3.246 ms / iter	1.516 ms / iter
After PR	3.240 ms / iter	1.461 ms / iter
Delta (%)	🟢 -0.18% (Faster)	🟢 -3.63% (Faster)

Without CuteDSL fusions

Metric	GPU-Bound Case (hidden_dim = 2048)	CPU-Bound Case (hidden_dim = 128)
Tokens	65,536	2,048
Before PR	7.353 ms / iter	1.654 ms / iter
After PR	3.398 ms / iter	1.377 ms / iter
Delta (%)	🟢 -53.79% (Faster)	🟢 -16.75% (Faster)

Fixes #2988

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

Please list the changes introduced in this PR:

• Change A
• Change B

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

This PR introduces three related MoE performance optimizations: (1) new C++ tex APIs for scaled activations (SReLU, SwiGLU, clamped SwiGLU) replacing multi-step PyTorch compositions; (2) ForwardScaledActivationGroupedLinear and BackwardScaledActivationGroupedLinear fused ops that combine the activation kernel with grouped quantization in a single dispatch, reducing CPU overhead; (3) precomputed tensor offsets (splits_to_offsets_multi) reused in the backward pass of both GroupedLinear and GroupedMLP, eliminating redundant kernel launches.

• New tex APIs & fusions: 12 pybind-exposed C++ template functions (scaled_swiglu, grouped_scaled_dswiglu, etc.) underpin both the standalone activation ops and the new fused ForwardScaledActivationGroupedLinear / BackwardScaledActivationGroupedLinear ops registered in the TE op fuser pipeline.
• Saved-tensor layout extension: _fuser_forward_grouped_tensor now saves [split_sizes, base_split_offsets, split_points, input_tensor_offsets, output_tensor_offsets, ...] so the backward can read precomputed offsets at fixed indices without re-running splits_to_offsets_multi; the split-quantize path pads with None to keep the layout contract uniform.
• Test coverage: fp8_current_scaling is added to the full grouped-MLP quantization test matrix and new assertions verify that ForwardScaledActivationGroupedLinear / BackwardScaledActivationGroupedLinear appear in forward/backward op lists when expected.

Confidence Score: 5/5

Safe to merge. The core fusion logic — tensor-offset reuse, grouped-scaled activation dispatch, and the new backward fused op — is internally consistent and confirmed correct by existing tests including the new CUTEDSL=0 run.

The forward and backward fused ops use precomputed tensor offsets that exactly match the layout recorded in forward, and the saved-tensor index contract is consistently documented across both the split-quantize and graph-safe paths. The C++ template helpers follow the existing TE extension pattern and pybind signatures match. The logic is stress-tested via the newly added NVTE_CUTEDSL_FUSED_GROUPED_MLP=0 test run that exercises the non-cuTe path end-to-end.

Files Needing Attention: tests/pytorch/test_grouped_mlp.py: the act_grouped_linear_fusion_expected flag couples both forward and backward fusion assertions with a single AND of fc1+fc2 support; if they differ the assertion could produce a false negative, worth a follow-up.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fused/forward_activation_grouped_linear.py	New fused op: ScaledActivation + GroupedLinear forward fusion. Correctly sequences grouped-scaled activation dispatch, input quantizer setup, and GroupedLinear's graph-safe forward. Saves input and scales for backward only when required. Logic and tensor-offset usage look correct.
transformer_engine/pytorch/ops/fused/backward_activation_grouped_linear.py	New fused op: FC1 + ScaledActivation backward fusion. Uses output_tensor_offsets from FC1's forward context (index 4) as tensor_offsets for the activation gradient quantization — correctly sized at fc1.out_features. Early-returns when linear_ctx.requires_grad is False. Return tuple layout matches TE fuser conventions.
transformer_engine/pytorch/csrc/extensions/activation.cpp	Adds C++ template helpers (scaled_activation_compute, scaled_dactivation_compute) and grouped variants, then exposes scaled_swiglu / scaled_srelu / grouped_scaled_* wrappers. Template design avoids duplication. The grouped_scaled_dactivation_helper returns a 3-tuple (quantized, dense, grad_scales) matching Python-side unpacking.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Saved-tensor layout extended with precomputed tensor offsets in _fuser_forward_grouped_tensor. Split-quantize backward path preserves layout compatibility by padding with None slots. New _fuser_backward_grouped_tensor refactors shared backward GEMM logic.
transformer_engine/pytorch/ops/basic/swiglu.py	ScaledSwiGLU and ScaledClampedQGeGLU forward/backward replaced with tex.scaled_* APIs. Scales are always saved (ctx.input_requires_grad is unconditionally True). No issues.
tests/pytorch/test_grouped_mlp.py	Adds fusion presence assertions and fp8_current_scaling to the full quantization test matrix. The act_grouped_linear_fusion_expected flag couples both forward and backward fusion checks with a single AND of fc1+fc2 support — could produce a false negative if they differ (see comment).
transformer_engine/pytorch/csrc/extensions.h	Adds declarations for 12 new scaled_* and grouped_scaled_* activation functions. Signatures match implementations in activation.cpp.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Pybinds the 12 new scaled_* activation C++ functions with correct argument names and default values.
transformer_engine/pytorch/ops/fused/init.py	Registers ForwardScaledActivationGroupedLinear and BackwardScaledActivationGroupedLinear fusions and exports act_grouped_linear_fusion_supported for test introspection.
transformer_engine/pytorch/module/grouped_linear.py	Precomputes input_tensor_offsets and output_tensor_offsets via splits_to_offsets_multi in forward and saves base_split_offsets and input_tensor_offsets for reuse in backward, avoiding redundant kernel calls.
transformer_engine/pytorch/ops/basic/activation.py	ScaledSReLU forward/backward now delegate to tex.scaled_srelu/scaled_dsrelu. Changes are minimal and correct.

Sequence Diagram

sequenceDiagram
    participant Input
    participant FwdFusion as ForwardScaledActivationGroupedLinear
    participant GroupQ as GroupQuantize (tex.grouped_scaled_*)
    participant FC2 as GroupedLinear FC2
    participant BwdFusion as BackwardScaledActivationGroupedLinear
    participant DActivation as _grouped_scaled_dactivation
    participant FC1Bwd as FC1._fuser_backward_grouped_tensor

    Note over Input,FC2: Forward pass
    Input->>FwdFusion: pre-activation tensor + scales
    FwdFusion->>GroupQ: compute + quantize activation (grouped_tensor_offsets[2])
    GroupQ-->>FwdFusion: grouped_x (quantized)
    FwdFusion->>FC2: grouped_x + split metadata
    FC2-->>FwdFusion: output + saved_tensors[0..4]
    FwdFusion-->>Input: output

    Note over BwdFusion,FC1Bwd: Backward pass (fused fc1 + act)
    BwdFusion->>DActivation: grad_output + saved input/scales
    Note right of DActivation: tensor_offsets = linear_ctx.saved_tensors[4] (output_tensor_offsets from FC1)
    DActivation-->>BwdFusion: grouped_dy (quantized), dense_dy, grad_scales
    BwdFusion->>FC1Bwd: dense_dy (wgrad) + grouped_dy (dgrad)
    FC1Bwd-->>BwdFusion: grad_input, grad_params, grad_extra_inputs
    BwdFusion-->>BwdFusion: return grad_input, [fc1_grads, ()], [fc1_extra, (grad_scales,)]

Loading

_{Reviews (10): Last reviewed commit: "better name" | Re-trigger Greptile}

[zhongbozhu]

There is a trade-off in creating a partially fused module, since we already have FC1 - ACT - FC2 fused module, We need to justify the value of creating a FC1-ACT fusion, instead of just adding features to the grouped_mlp instead.

[vthumbe1503]

This essentially enables the infrastructure to enable (Act + GroupQuant from FC2) fusion, and enables to add fused kernels if possible in the future. cc: @timmoon10

[vthumbe1503]

The other important thought i had in mind is the fact that ScaledActivation today doesnt take in m_splits and it is going to be a lot of upstream disruption to allow for that. And at the same time ScaledActivation can be used after Dense Layers as well and not necessarily after a GroupedLinear layer(and so it might not always need m_splits).

Allowing for this fusion, we allow the m_splits information to be also consumed in the scaled_activation + grouped quantization fusion. Right now the activation kernel that we are using isnt even using the m_splits, but for paged stashing optimization we might need that and we can potentially add a new kernel for that in the future.

[vthumbe1503]

/te-ci pytorch

[vthumbe1503]

/te-ci pytorch