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. However the same is done 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 adds three related optimizations to the MoE grouped-linear stack: (1) new tex C++ APIs for scaled activations (srelu, swiglu, clamped-swiglu) and their fused grouped-quantize variants that replace manual Python-level scale multiplication and interleave/deinterleave logic; (2) pairwise ForwardScaledActivationGroupedLinear / BackwardScaledActivationGroupedLinear fused ops that collapse the activation→grouped-quantize→linear boundary into a single kernel dispatch, reducing CPU overhead; and (3) precomputed tensor offsets (splits_to_offsets_multi) that are saved in the forward context and reused in backward, eliminating 2–4 runtime multiply-and-quantize kernel calls per step in both GroupedLinear and GroupedMLP.

• Scaled-activation C++ APIs: Templated scaled_activation_compute / grouped_scaled_dactivation_helper helpers unify all variants; backward returns a 3-tuple (grouped-quantized, dense, grad_scales) consumed correctly by the Python fused backward.
• Activation+GroupedLinear fusions: Forward fusion covers (ScaledActivation → fc2) and backward fusion covers (fc1 ← ScaledActivation); tensor-offset bookkeeping (output_tensor_offsets at saved_tensors[4]) is reused across both boundaries.
• Precomputed offsets: splits_to_offsets_multi now computes fc1-input, fc1-output, fc2-input, and fc2-output offsets in one call for GroupedMLP, removing runtime base_split_offsets × stride multiplications from both forward and backward.

Confidence Score: 5/5

Safe to merge; fused forward and backward paths are correctly wired, tensor-offset bookkeeping is consistent across all call sites.

The gradient math is correct end-to-end. The saved-tensor layout at index [4] (output_tensor_offsets) is consistently written in forward and read in backward across GroupedLinear, ForwardScaledActivationGroupedLinear, and BackwardScaledActivationGroupedLinear. The unreachable branch does not affect correctness.

Files Needing Attention: backward_activation_grouped_linear.py lines 115–121 contain a dead branch with an incorrect gradient-routing assumption that could matter if the op framework semantics change.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fused/forward_activation_grouped_linear.py	New fused forward op that applies scaled activation + grouped quantize + grouped linear in one step, reducing CPU kernel launch overhead; logic and offset bookkeeping are correct.
transformer_engine/pytorch/ops/fused/backward_activation_grouped_linear.py	New fused backward op; the if not activation_ctx.requires_grad branch (lines 115–121) is unreachable under TE's op-framework semantics and constitutes dead code, though it introduces no runtime error.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Refactors _fuser_forward/backward_grouped_tensor to accept pre-built GroupedTensorStorage inputs; adds precomputed tensor offsets to saved-tensor layout; saves 2-4 device kernel calls per step.
transformer_engine/pytorch/ops/basic/swiglu.py	Removes manual interleave/de-interleave logic; delegates to tex.scaled_swiglu/scaled_clamped_swiglu; fixes a pre-existing bug where scales were not saved when only extra_input_requires_grad was True.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Precomputes fc1_out_tensor_offsets alongside existing offsets; eliminates runtime multiply (base_split_offsets × stride) in both fc1 and fc2 backward paths.
tests/pytorch/test_grouped_mlp.py	Adds assertions for forward/backward activation-grouped-linear fusions; fusion-expected predicate ANDs fc1 and fc2 support checks, which can suppress a true positive when only one boundary's support differs.
transformer_engine/pytorch/csrc/extensions/activation.cpp	Adds templated C++ helpers for scaled_activation_compute and grouped_scaled_dactivation_helper; backward variant returns a 3-tuple (grouped-quantized, dense, grad_scales) as expected by Python callers.
transformer_engine/pytorch/module/grouped_linear.py	Precomputes input and output tensor offsets together with split_sizes via splits_to_offsets_multi; passes them directly to group_quantize calls in forward and backward.

Sequence Diagram

sequenceDiagram
    participant Input
    participant ScaledActivation
    participant GroupedQuantize as Grouped Quantize
    participant FC2 as GroupedLinear (fc2)
    participant Output

    Note over Input,Output: Forward Pass (ForwardScaledActivationGroupedLinear)
    Input->>ScaledActivation: input_ [tokens, 2H or H]
    ScaledActivation->>GroupedQuantize: activation output [tokens, H]
    GroupedQuantize->>FC2: grouped_x (quantized)
    FC2->>Output: out [tokens, out_features]

    Note over Output,Input: Backward Pass (BackwardScaledActivationGroupedLinear)
    Output-->>FC2: grad_output [tokens, fc1.out_features]
    FC2-->>GroupedQuantize: grouped_dy (quantized dactivation input)
    GroupedQuantize-->>ScaledActivation: dense_dy [tokens, fc1.out_features]
    ScaledActivation-->>Input: grad_input [tokens, fc1.in_features]

Loading

_{Reviews (8): Last reviewed commit: "Merge branch 'main' into grouped_linear_..." | 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