[Common][PyTorch] Add SiTU-GLU activation

[harryzhou2000]

Description

Add native and row-scaled SiTU-GLU operations and connect them to the fused grouped-MLP path when the installed cuDNN frontend exposes SiTU-GLU support.

SiTU-GLU is the soft-capped gated activation used by Kimi K3. For gate input a, up input b, and positive cap parameters beta1, beta2, this PR computes

gate(a) = beta1 * tanh(a / beta1) * sigmoid(a)
up(b)   = beta2 * tanh(b / beta2)
y       = gate(a) * up(b)

The backward kernels use

d gate / da = (1 - tanh(a / beta1)^2) * sigmoid(a)
              + beta1 * tanh(a / beta1) * sigmoid(a) * (1 - sigmoid(a))
d up / db   = 1 - tanh(b / beta2)^2

The default values beta1=4 and beta2=25 follow the Kimi K3 technical report.

The grouped-GEMM integration depends on the cuDNN frontend support proposed at NVIDIA/cudnn-frontend#645. Compatibility is detected from the forward and backward wrapper signatures instead of from a package-version threshold. An older or unpatched cuDNN frontend therefore keeps the correct native GroupedLinear -> ScaledSiTUGLU -> GroupedLinear fallback, while a frontend that exposes both situ_beta1 and situ_beta2 parameters enables the fused path.

Type of change

• Documentation change
• Bug fix
• New feature
• Breaking change
• Infra/Build change
• Code refactoring

Changes

• Add public C APIs and PyTorch bindings for SiTU-GLU forward/backward.
• Add te.ops.SiTUGLU with contiguous and block-interleaved layouts and existing FP8/MXFP8 quantizer integration.
• Add row-scaled te.ops.ScaledSiTUGLU, including optional scale gradients and a constructor-compatible recomputation flag that explicitly rejects unsupported True requests.
• Reject activation_recompute_in_mlp=True for scaled SwiGLU, SiTU-GLU, and clamped QGeGLU instead of silently ignoring it in the fused grouped MLP. Scaled SReLU remains supported.
• Reuse TE's vectorized, FP8, and MXFP8 gated kernels for native execution; no PyTorch-composed activation or intermediate gate/up tensors are required.
• Add ScaledSiTUGLU to grouped-MLP dimension validation and pass situglu/dsituglu plus both cap parameters to cuDNN frontend forward/backward wrappers.
• Cache a feature probe that requires both cap parameters in both grouped wrapper signatures before registering the fusion.
• Add common C++ references and PyTorch forward/backward, quantization, interleave, fallback, feature-detection, and fused grouped-MLP tests.

Performance and compatibility

• Native BF16/FP16/FP32 execution is a single TE gated activation kernel. FP8/MXFP8 quantization remains fused into the existing gated kernels.
• Row scaling is handled by TE's existing scaled-activation kernels rather than a separate PyTorch multiply.
• With [CuTeDSL] Add grouped SiTU-GLU activation cudnn-frontend#645, SiTU-GLU stays inside the grouped FC1 epilogue and dSiTU-GLU inside the grouped backward kernel, so the grouped path does not add a standalone activation launch or materialized FC1 activation.
• The signature probe is cached and only runs during fusion setup; it is not in the per-iteration GPU path.
• The stock 26.06 container stack (nvidia-cudnn-frontend==1.26.0, nvidia-cutlass-dsl==4.5.0) correctly reports the feature as unavailable and uses the native fallback.
• A private build of [CuTeDSL] Add grouped SiTU-GLU activation cudnn-frontend#645 (package version 1.27.0) with the same CuTe DSL 4.5.0 runtime reports the feature as available and passes fused MXFP8 forward/backward tests.

Validation

Run on one NVIDIA B300 GPU in the 26.06 container using an isolated private venv:

• TE build for NVTE_CUDA_ARCHS="100;103a;" passed. Optional NCCL-EP was disabled because the separately reused NCCL-EP headers predate upstream TE's current API; no SiTU code depends on NCCL-EP.
• 24 passed: focused native/scaled SiTU-GLU tests with stock cuDNN frontend, including FP32/FP16/BF16, interleaving, invalid parameters, MXFP8 output/input-gradient quantization, scale gradients, unsupported recomputation rejection, feature detection, and grouped fallback.
• 2 passed: grouped MXFP8 fallback with (beta1, beta2)=(4,25) and (2,8) on stock cuDNN frontend 1.26.0.
• 2 passed: fused grouped MXFP8 forward/backward with the same two parameter pairs using [CuTeDSL] Add grouped SiTU-GLU activation cudnn-frontend#645.
• 244 passed, 206 skipped: existing SwiGLU-focused native, quantized, grouped-MLP, and MCore-integration regression selection.
• All targeted pre-commit hooks passed (Black 24.4.2, clang-format 18.1.6, whitespace, merge-conflict, large-file, and Python 3.10 checks).

Checklist

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented hard-to-understand areas
• My changes generate no new warnings
• I have added tests that prove the feature works
• New and existing focused unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

The PR adds native and row-scaled SiTU-GLU operations and conditionally integrates them with the fused grouped-MLP path.

• Adds C/CUDA kernels, public APIs, PyTorch bindings, and quantized execution support.
• Adds cuDNN frontend feature detection and native fallback behavior.
• Rejects unsupported recomputation for scaled GLUs while retaining supported ScaledSReLU recomputation.
• Adds forward, backward, quantization, layout, fallback, and grouped-fusion tests.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the previously reported recomputation no-op is resolved by rejecting unsupported scaled-GLU requests while leaving supported ScaledSReLU recomputation intact.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/basic/swiglu.py	Adds SiTU-GLU operations and correctly rejects unsupported scaled-GLU recomputation while preserving the constructor-compatible false default.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Adds feature-gated SiTU-GLU grouped-MLP fusion and retains SReLU as the sole supported activation-recomputation path.
transformer_engine/common/activation/scaled_activation.cu	Extends shared scaled gated kernels with SiTU-GLU forward and derivative handling.
transformer_engine/common/activation/swiglu.cu	Adds validated native SiTU-GLU forward and backward C API implementations.
transformer_engine/pytorch/csrc/extensions/activation.cpp	Connects the new native and scaled SiTU-GLU APIs to PyTorch tensors and quantizers.
tests/pytorch/test_fusible_ops.py	Covers SiTU-GLU numerics, quantization, scaling, parameter validation, and recomputation rejection.
tests/pytorch/test_grouped_mlp.py	Covers feature detection, native fallback, and real cuDNN fused grouped-MLP execution.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[GroupedLinear FC1] --> B{cuDNN frontend supports SiTU-GLU parameters?}
  B -->|Yes| C[Fused grouped SiTU-GLU path]
  B -->|No| D[Native ScaledSiTUGLU fallback]
  C --> E[GroupedLinear FC2]
  D --> E

Loading

_{Reviews (8): Last reviewed commit: "[Common] Fix scaled activation scale-gra..." | Re-trigger Greptile}

[harryzhou2000]

Update in 730c0c15: complete NVFP4 SiTU-GLU integration with the cuDNN GLU-Hadamard path from NVIDIA/cudnn-frontend#645.

TE now feature-detects situ_beta1 and situ_beta2 on grouped_gemm_glu_hadamard_wrapper_sm100. NVFP4 SiTU uses fused grouped GEMM + SiTU-GLU + post-RHT amax when available, while an older cuDNN frontend keeps the existing regular grouped GEMM + SiTU path. SwiGLU and SReLU behavior is unchanged.

The test update also makes fused-path validation strict: it traces the real cuDNN wrappers, verifies forward/backward activation names and generated tensor dtypes, and does not accept silent native fallback when fusion is required.

Validation on NVIDIA B200:

• 3 focused TE tests passed: Hadamard feature detection and real NVFP4 SiTU-Hadamard execution for (4,25) and (2,8).
• A real MCore MoE forward/backward passed with TE native ScaledSiTUGLU and GroupedMLP_CuTeGEMMGLU.
• Wrapper tracing observed only the cuDNN GLU-Hadamard forward call (situglu, FP32 alpha) and regular dGLU backward call (dsituglu, FP32 alpha/beta); the regular grouped-GLU forward wrapper was not called.
• The corresponding MXFP8 strict integration still passes through regular grouped GLU/dGLU with BF16 alpha/beta.

[timmoon10]

/te-ci

[harryzhou2000]

Addressed the gated-derivative review feedback in 15dd10f.

The native FP8, MXFP8 colwise/rowwise, vectorized, and row-scaled implementations now use one contract:

• act_elt = act(x1), gate_elt = gate(x2)
• dact_elt = dact/dx1, dgate_elt = dgate/dx2
• dx1 = grad * dact_elt * gate_elt
• dx2 = grad * act_elt * dgate_elt

For ordinary GLUs, dgate_elt is 1. For clamped GLUs it is the 0/1 clamp derivative, and for SiTU-GLU it is the smooth SiTU up-branch derivative. This removes the SiTU-only conditional composition and makes the row-scaled helper follow the same semantics as the FP8/MXFP8 kernels.

Focused validation on B300 with PyTorch 26.06, cuDNN frontend 1.26.0 source, and CuTe DSL 4.6.2:

• Editable TE build for sm100 and sm103a: passed
• Focused fusible-op matrix: 175 passed, 42 capability-based skips, 0 failures
• Covered native/vectorized SiTU-GLU, MXFP8 SiTU-GLU, clamped SwiGLU quantization variants, scaled SiTU/SwiGLU, and scaled clamped forward/backward
• Identical rebuild after cache normalization: 33 direct ccache hits, 0 misses

Related cuDNN-frontend patch for the SiTU-GLU Hadamard API contracts: NVIDIA/cudnn-frontend#670

[timmoon10]

LGTM, pending CI

[timmoon10]

/te-ci

[vthumbe1503]

One of the optimizations that we did for swiglu and clamped_swiglu was to do the exponential computations shared by both the activation and dactivation functions once, as you can see in this gated kernel.

For example in this case, situ and dsitu can share the tanh computation being done once and use it. Ideally compiler should be smart enough to do those shared computations once. But we did find it wasnt the case for swiglu and clamped_swiglu in terms of reusing exp computations.

Not a blocker for this PR, but we should definitely see if the perf of this kernel becomes better if do the same optimization
cc: @ptrendx @Oleg-Goncharov