Activation + GroupedLinear Fusion for MOE and other MOE optimizations

[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 introduces a fused ScaledActivation + GroupedLinear forward/backward operation for MoE (Mixture of Experts) inference and training, alongside several supporting refactors: splitting the now-overloaded _fuser_forward_grouped_tensor / _fuser_backward_grouped_tensor into a "build grouped storage" outer shell and a "run GEMMs" inner method, pre-computing all tensor-offset arrays via splits_to_offsets_multi (replacing arithmetic on base_split_offsets at use sites), and adding the new C++ scaled_* and grouped_scaled_* activation kernels with pybind bindings.

• New fusion pair: ForwardScaledActivationGroupedLinear / BackwardScaledActivationGroupedLinear fuse the per-row-scaled activation (SReLU / SwiGLU / ClampedSwiGLU) with the subsequent grouped GEMM into a single pass, eliminating a separate grouped-quantize step.
• Offset pre-computation refactor: GroupedLinear._fuser_forward_graph_safe and _GroupedMLP_CuTeGEMMBase now call splits_to_offsets_multi upfront and save the resulting input_tensor_offsets / output_tensor_offsets tensors, replacing the runtime base_split_offsets * stride multiplications in both forward and backward.
• Bug fix in _ScaledGLU: scales are now saved for backward whenever input_requires_grad OR extra_input_requires_grad, fixing a case where the scale gradient could not be computed when only the scales needed gradient.

Confidence Score: 4/5

Safe to merge for normal training configurations; the two new fused ops have no known correctness issues for the common path where both linear and activation contexts require gradients.

The offset pre-computation refactor and the new fused activation+linear ops are structurally sound. The unfused fallback paths are preserved and the saved-tensor layout contract is consistently maintained across all callers. The only issues found are in low-traffic branches: input_requires_grad is always set True (wastes computation but not incorrect in the fully-fused path), and the activation scales tensor is not released after the backward runs.

Files Needing Attention: forward_activation_grouped_linear.py and backward_activation_grouped_linear.py warrant a careful second look on the early-return branch (not linear_ctx.requires_grad) and the tensor-clear coverage.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fused/forward_activation_grouped_linear.py	New fused forward op: activation_ctx.input_requires_grad is hardcoded to True; minor performance waste in the not-linear-requires-grad backward branch.
transformer_engine/pytorch/ops/fused/backward_activation_grouped_linear.py	New fused backward op: only saved_tensors[0] (input_) is cleared after use; scales tensor at saved_tensors[1] is not cleared, unlike the unfused _ScaledGLU path.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Refactors _fuser_forward_grouped_tensor into graph_safe outer + inner methods; expands saved-tensor prefix from 3 to 5 entries; both paths now use pre-computed tensor offsets consistently.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Adds fc1_out_tensor_offsets to splits_to_offsets_multi call and saves 4 offset tensors (7-entry prefix) instead of 3; replaces base_split_offsets arithmetic at all backward use sites.
transformer_engine/pytorch/ops/basic/swiglu.py	Moves GLU interleave reshaping into C++ scaled_swiglu/scaled_dswiglu kernels; fixes scale save condition to include extra_input_requires_grad.
transformer_engine/pytorch/ops/basic/activation.py	Extracts _ScaledUnary ABC from ScaledSReLU; routes forward/backward through new tex.scaled_srelu / tex.scaled_dsrelu kernels; exposes _scaled_unary_forward/_backward for fusion consumers.
transformer_engine/pytorch/csrc/extensions/activation.cpp	Adds scaled_activation_compute / scaled_dactivation_compute templates and grouped variants; correct shape checks and GIL-release for all new kernels.

Sequence Diagram

sequenceDiagram
    participant FW as ForwardScaledActivationGroupedLinear
    participant ACT as _grouped_scaled_activation (C++ kernel)
    participant LIN as GroupedLinear._fuser_forward_grouped_tensor
    participant BW as BackwardScaledActivationGroupedLinear
    participant DACT as _grouped_scaled_dactivation (C++ kernel)
    participant DLIN as GroupedLinear._fuser_backward_grouped_tensor

    note over FW,LIN: Forward pass
    FW->>FW: splits_to_offsets_multi → [split_pts, base_offs, in_offs, out_offs]
    FW->>ACT: input_, scales, input_quantizer, split_sizes, in_offs
    ACT-->>FW: grouped_x (quantized GroupedTensorStorage)
    FW->>FW: activation_ctx.save_for_backward(input_, scales)
    FW->>LIN: grouped_x, split_sizes, out_offs, ...
    LIN-->>FW: out tensor + tensors_to_save (5-entry prefix layout)
    FW->>FW: linear.fuser_forward_save_ctx(linear_ctx, tensors_to_save)

    note over BW,DLIN: Backward pass
    BW->>BW: read activation_ctx.saved_tensors → input_, scales
    BW->>BW: read linear_ctx.saved_tensors[4] → output_tensor_offsets
    BW->>DACT: grad_output, input_, scales, grad_output_quantizer, out_offs
    DACT-->>BW: grouped_dy (quantized), dense_dy, grad_scales
    BW->>DLIN: grouped_dy, dense_dy as grad_output
    DLIN-->>BW: grad_input, grad_params, grad_extra_inputs
    BW-->>BW: return grad_input, weight_grads, scale_grads

Loading

_{Reviews (7): 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