[PyTorch]Fused ScaledActivation Kernel in TE Ops and reducing CPU overheads for MOE

[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 and used directly in the corresponding scaled te ops for activations

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.

Also strengthens test_grouped_mlp assertions. We expect fusion even for nvfp4_rht + scale_swiglu + glu_interleave_size = 32 case.

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.308 ms / iter
Delta (%)	🟢 -53.79% (Faster)	🟢 -20.91% (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 fused C++ kernels for scaled activations (SReLU, SwiGLU, clamped SwiGLU), replaces the multi-step Python-level arithmetic in ScaledSReLU/ScaledSwiGLU/ScaledClampedQGeGLU with single tex.scaled_* calls, and precomputes all grouped-tensor byte offsets in the forward pass so the backward can reuse them instead of recomputing.

• Fused scaled-activation kernels: New tex.scaled_srelu, tex.scaled_swiglu, tex.scaled_clamped_swiglu (and their d* backward counterparts) are wired through activation.cpp and swiglu.py; a shared _ScaledUnary ABC and updated _ScaledGLU route forward/backward through these kernels while the GLU-interleaving reshape is now handled entirely inside the C++ kernel.
• Precomputed tensor offsets: Both GroupedLinear and _GroupedMLP_CuTeGEMMBase now call splits_to_offsets_multi once per forward to bulk-allocate all per-group byte-offset tensors (input_tensor_offsets, output_tensor_offsets, fc1_out_tensor_offsets, etc.) and save them for backward, eliminating several redundant offset-computation kernel launches.

Confidence Score: 5/5

Safe to merge — the changes are internally consistent refactors with no correctness regressions identified.

All forward/backward offset index adjustments are consistent across forward save and backward unpack sites. The new fused C++ kernels preserve the mathematical semantics of the old multi-step Python paths (scale multiply, activation, interleave reshape). The shared-expert (num_groups==1) path correctly sets new offset tensors to None, and the branch guard ensures the nvfp4 dgrad path that consumes fc1_x_tensor_offsets is never reached for that case. The only dead-code guard (if not ctx.input_requires_grad) is harmless.

Files Needing Attention: No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/activation.cpp	Adds C++ template helpers scaled_activation_compute and scaled_dactivation_compute plus maybe_quantize; implements all six new scaled_* / scaled_d* functions. Logic looks correct — grad_input is always allocated (no skip path when compute_scale_grad=false only), which is consistent with the Python callers always requiring it.
transformer_engine/pytorch/csrc/extensions.h	Adds declarations for the six new scaled-activation public C++ functions; straightforward header-only change.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Registers the six new functions with pybind11; default argument values match those in the C++ signatures.
transformer_engine/pytorch/ops/basic/activation.py	Introduces _ScaledUnary ABC; ScaledSReLU now delegates to tex.scaled_srelu/tex.scaled_dsrelu. The if not ctx.input_requires_grad: grad_input = None guard in fuser_backward is dead code (the forward hardcodes ctx.input_requires_grad = True), but harmless.
transformer_engine/pytorch/ops/basic/swiglu.py	Replaces manual de-interleave + tex.swiglu + pointwise multiply with single tex.scaled_swiglu/tex.scaled_clamped_swiglu calls; scales-saving bug fix (ctx.input_requires_grad or ctx.extra_input_requires_grad) is correct, and interleaving is now fully handled in the C++ kernel.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Passes precomputed input_tensor_offsets/output_tensor_offsets to all group-quantize and GroupedTensorStorage calls; saved-tensor layout extended from 3→5 metadata slots, and all read-back indices updated consistently in both _fuser_backward_split_quantize and _fuser_backward_grouped_tensor. CPU-offload slot offset also correctly bumped from 3/4 to 5/6.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Adds fc1_out_tensor_offsets to the bulk-allocated offset set; all four precomputed offsets are saved/restored via the extended 7-element metadata prefix; backward uses these offsets directly instead of recomputing. The shared-expert (num_groups==1) path correctly sets new offsets to None, and the use_single_group_dense_dgrad branch guards ensure the nvfp4 path (which uses fc1_x_tensor_offsets) is never reached for num_groups==1.
transformer_engine/pytorch/module/grouped_linear.py	Uses splits_to_offsets_multi with strides=[1, in_features, out_features] to replace the old splits_to_offsets + multiply pattern; offsets saved alongside weights and correctly unpacked in backward.
tests/pytorch/test_grouped_mlp.py	Adds NVTE_GROUPED_LINEAR_SINGLE_PARAM env-var guards to four test methods and expands the nvfp4_rht BF16 fusion-expected condition to cover all non-GLU activations and GLU with glu_interleave_size==32.

Sequence Diagram

sequenceDiagram
    participant PY as Python (fuser_forward)
    participant TEX as tex (C++ ext)
    participant KERNEL as CUDA Kernel

    Note over PY,KERNEL: Forward Pass
    PY->>TEX: "splits_to_offsets_multi(split_sizes, strides=[...])"
    TEX-->>PY: split_sizes, [base_offsets, in_offsets, out_offsets, ...]
    PY->>TEX: "group_quantize(x, quantizer, num_groups, split_sizes, tensor_offsets=in_offsets)"
    TEX->>KERNEL: fused quantize+group
    KERNEL-->>PY: grouped_x
    PY->>TEX: scaled_swiglu(input, scales, quantizer, glu_interleave_size)
    TEX->>KERNEL: nvte_scaled_swiglu (fused act+scale)
    KERNEL-->>PY: activation_out
    Note over PY: save [split_sizes, base_offsets, split_points, in_offsets, out_offsets, ...]

    Note over PY,KERNEL: Backward Pass (reuses saved offsets)
    PY->>TEX: "group_quantize(dy, quantizer, N, split_sizes, tensor_offsets=out_offsets)"
    TEX->>KERNEL: fused quantize+group
    KERNEL-->>PY: grouped_dy
    PY->>TEX: scaled_dswiglu(grad, input, scales, quantizer, glu_interleave_size, compute_scale_grad)
    TEX->>KERNEL: nvte_scaled_dswiglu (fused grad_input + grad_scales)
    KERNEL-->>PY: grad_input, grad_scales

Loading

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

[vthumbe1503]

/te-ci pytorch

[timmoon10]

LGTM

[vthumbe1503]

Pipeline: 61451302