Activation + GroupedLinear Fusion for MOE

[vthumbe1503]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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 a fused ScaledActivation + GroupedLinear kernel path for MoE models, allowing the activation (SwiGLU, clamped SwiGLU, SReLU) and the subsequent grouped quantize + GEMM to execute as a single kernel dispatch in both forward and backward. It also refactors GroupedLinear to separate grouped-tensor-input construction from GEMM execution, enabling external callers to supply a pre-quantized GroupedTensorStorage.

• New C++ kernels (grouped_scaled_swiglu, grouped_scaled_srelu, and their backward variants) fuse per-row scale multiplication with activation and grouped quantization into one pass; Python-side ForwardScaledActivationGroupedLinear / BackwardScaledActivationGroupedLinear wire these into the TE fuser framework.
• GroupedLinear._fuser_forward_grouped_tensor is renamed _fuser_forward_graph_safe and a new _fuser_forward_grouped_tensor accepting a pre-built GroupedTensorStorage is added, enabling the fused activation op to skip redundant quantize steps when passing the activation output directly into the GEMM.
• ScaledSReLU is refactored into an abstract _ScaledUnary base class to share the unfused forward/backward logic across unary scaled activations.

Confidence Score: 4/5

Safe to merge for standard full-training runs; the new fused path has an untested crash for partially-frozen MoE configurations.

The fused forward and backward kernels are correct for the common training case where all ops have requires_grad=True. However, BackwardScaledActivationGroupedLinear.fuser_backward unconditionally reads linear_ctx.saved_tensors[0] regardless of whether fc1's context actually saved anything. When fc1 weights are frozen but the per-row scales are trainable, the TE fuser still enters the fused backward, hits a NoneType subscript error, and crashes. This is a real failure mode in fine-tuning scenarios and is not covered by any test.

transformer_engine/pytorch/ops/fused/backward_activation_grouped_linear.py — the fuser_backward method needs a guard around the linear_ctx.saved_tensors access when linear_ctx.requires_grad is False.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fused/backward_activation_grouped_linear.py	New fused backward for (fc1, activation) pair; crashes when fc1's context has requires_grad=False while activation's does not, because saved_tensors is None.
transformer_engine/pytorch/ops/fused/forward_activation_grouped_linear.py	New forward fused op (activation + fc2 linear); correctly fuses scaled activation + quantize into a single grouped kernel call and saves per-activation context for backward.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Renames _fuser_forward_grouped_tensor to _fuser_forward_graph_safe and introduces a new _fuser_forward_grouped_tensor that accepts a pre-built GroupedTensorStorage, enabling the fused activation path to reuse the GEMM logic.
transformer_engine/pytorch/csrc/extensions/activation.cpp	Adds grouped_scaled_activation_helper / grouped_scaled_dactivation_helper templates and six concrete dispatch functions for SwiGLU, clamped SwiGLU, and SReLU (forward + backward).
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Registers the six new grouped-scaled activation kernels in the pybind11 module with correct argument names and defaults.
transformer_engine/pytorch/csrc/extensions.h	Adds C++ declarations for the six new grouped scaled activation functions.
transformer_engine/pytorch/ops/fused/init.py	Exports new forward/backward fused ops and registers them with the TE forward and backward fusion hooks.
tests/pytorch/test_grouped_mlp.py	Adds test_grouped_mlp_act_grouped_linear_fusion, guards single-weight tests behind the env-var flag, and widens the act_grouped_linear fusion assertion to cover both fc1 and fc2 (fc1 check is spurious for the fc2-side fusion).

Sequence Diagram

sequenceDiagram
    participant Fuser
    participant ForwardSAGL as ForwardScaledActivationGroupedLinear
    participant tex as tex (C++ kernels)
    participant GroupedLinear as GroupedLinear._fuser_forward_grouped_tensor

    Note over Fuser: Forward pass (activation→fc2 fused)
    Fuser->>ForwardSAGL: fuser_forward(input_, scales, split_sizes)
    ForwardSAGL->>tex: grouped_scaled_swiglu/srelu(x, scales, quantizer, split_sizes)
    tex-->>ForwardSAGL: grouped_x (quantized activation output)
    ForwardSAGL->>GroupedLinear: _fuser_forward_grouped_tensor(grouped_x, ...)
    GroupedLinear-->>ForwardSAGL: (out, tensors_to_save)
    ForwardSAGL->>ForwardSAGL: save_for_backward(input_, scales) on activation_ctx
    ForwardSAGL-->>Fuser: out

    Note over Fuser: Backward pass (fc1 backward fused with activation backward)
    Fuser->>Fuser: GroupedLinear(fc2).fuser_backward(dL/dfc2_out)
    Note over Fuser: Returns dL/d(activation_out * scales) as grad_output

    participant BackwardSAGL as BackwardScaledActivationGroupedLinear
    Fuser->>BackwardSAGL: fuser_backward(dL/d_activation_out)
    BackwardSAGL->>tex: grouped_scaled_dswiglu/dsrelu(grad, input_, scales, quantizer)
    tex-->>BackwardSAGL: (grouped_dy, dense_dy, grad_scales)
    BackwardSAGL->>GroupedLinear: _fuser_backward_grouped_tensor(fc1_ctx, dense_dy, grouped_dy)
    GroupedLinear-->>BackwardSAGL: (dL/dfc1_input, wgrad_fc1, ...)
    BackwardSAGL-->>Fuser: (grad_input, grad_params, grad_extra_inputs)

Loading

_{Reviews (2): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

Crash when fc1 is frozen but activation scales are trainable

linear_ctx.saved_tensors is None whenever linear_ctx.requires_grad=False, because fuser_forward_save_ctx returns early (if not requires_grad[0]: return) without ever calling save_for_backward. The backward loop only skips this fused op when all of its basic-op contexts have requires_grad=False; when fc1 is frozen (so first_op_requiring_backward = idx_of_activation) but the per-row scales are trainable, fc1's context has requires_grad=False while the activation context has requires_grad=True. The loop therefore enters fuser_backward, which unconditionally does linear_ctx.saved_tensors[0], raising TypeError: 'NoneType' object is not subscriptable.

A concrete path: MoE fine-tuning where fc1 weights are frozen, the module input is a non-leaf tensor, and the scales (activation extra input) carry requires_grad=True. The fix is either to guard the access with if linear_ctx.requires_grad: and skip the linear backward when fc1 doesn't need it, or to save split_sizes directly on activation_ctx so the offset computation does not depend on linear_ctx.saved_tensors.