[PyTorch] Add SiLUGate fusible op

docs/api/pytorch.rst

@@ -259,6 +259,8 @@ Operation fuser
 
 .. autoapiclass:: transformer_engine.pytorch.ops.RMSNorm
 
+.. autoapiclass:: transformer_engine.pytorch.ops.SiLUGate
+
 .. autoapiclass:: transformer_engine.pytorch.ops.SReGLU
 
 .. autoapiclass:: transformer_engine.pytorch.ops.SReLU

tests/pytorch/test_fusible_ops.py

@@ -2332,6 +2332,164 @@ def test_rmsnorm(
         torch.testing.assert_close(dx_test, x_ref.grad, **tols)
         torch.testing.assert_close(dw_test, w_ref.grad, **tols)
 
+    @pytest.mark.parametrize("in_shape", ((1, 37, 16, 128), (6, 16, 64)))
+    @pytest.mark.parametrize("dtype", _dtypes)
+    def test_silu_gate(
+        self,
+        *,
+        in_shape: Iterable[int],
+        dtype: torch.dtype,
+        device: torch.device = "cuda",
+    ) -> None:
+        """SiLU gate"""
+
+        # Random data
+        x_ref, x_test = make_reference_and_test_tensors(
+            in_shape,
+            test_dtype=dtype,
+            test_device=device,
+        )
+        # Signed gates cover SiLU's negative side, where its derivative changes sign
+        g_ref, g_test = make_reference_and_test_tensors(
+            in_shape,
+            min=-6.0,
+            max=4.0,
+            test_dtype=dtype,
+            test_device=device,
+        )
+        dy_ref, dy_test = make_reference_and_test_tensors(
+            in_shape,
+            test_dtype=dtype,
+            test_device=device,
+            requires_grad=False,
+        )
+
+        # Plain PyTorch implementation. The op computes in FP32 and rounds
+        # the result to the input dtype, so the reference does the same.
+        y_ref = (x_ref.float() * torch.nn.functional.silu(g_ref.float())).to(dtype)
+        y_ref.backward(dy_ref.to(dtype))
+        dx_ref = x_ref.grad.to(dtype)
+        dg_ref = g_ref.grad.to(dtype)
+
+        # Implementation with fusible operation
+        op = te_ops.SiLUGate()
+        y_test = op(x_test, g_test)
+        y_test.backward(dy_test)
+
+        # Expected numerical error
+        tols = dtype_tols(dtype)
+
+        # Check results
+        y_test = y_test.to(dtype=torch.float64, device="cpu")
+        dx_test = x_test.grad.to(dtype=torch.float64, device="cpu")
+        dg_test = g_test.grad.to(dtype=torch.float64, device="cpu")
+        torch.testing.assert_close(y_test, y_ref.to(dtype=torch.float64, device="cpu"), **tols)
+        torch.testing.assert_close(dx_test, dx_ref.to(dtype=torch.float64, device="cpu"), **tols)
+        torch.testing.assert_close(dg_test, dg_ref.to(dtype=torch.float64, device="cpu"), **tols)
+
+    @pytest.mark.parametrize("in_shape", ((1, 37, 16, 128), (6, 16, 64)))
+    @pytest.mark.parametrize("dtype", (torch.float32, torch.bfloat16))
+    @pytest.mark.parametrize("zero_centered_gamma", (False, True))
+    @pytest.mark.parametrize("gate_layout", ("contiguous", "strided"))
+    def test_rmsnorm_silu_gate(
+        self,
+        *,
+        in_shape: Iterable[int],
+        dtype: torch.dtype,
+        device: torch.device = "cuda",
+        eps: float = 1e-6,
+        zero_centered_gamma: bool,
+        gate_layout: str,
+    ) -> None:
+        """RMSNorm with SiLU gate"""
+
+        # Skip unsupported BF16
+        if dtype == torch.bfloat16 and not is_bf16_available():
+            pytest.skip("BF16 requires sm_80 or higher")
+
+        # Make input and weight shapes consistent
+        in_shape = list(in_shape)
+        weight_shape = (in_shape[-1],)
+
+        # Random data
+        x_ref, x_test = make_reference_and_test_tensors(
+            in_shape,
+            test_dtype=dtype,
+            test_device=device,
+        )
+        w_ref, w_test = make_reference_and_test_tensors(
+            weight_shape,
+            test_dtype=dtype,
+            test_device=device,
+        )
+        dy_ref, dy_test = make_reference_and_test_tensors(
+            in_shape,
+            test_dtype=dtype,
+            test_device=device,
+            requires_grad=False,
+        )
+
+        # Gate tensor, optionally a non-contiguous view of a larger tensor
+        if gate_layout == "contiguous":
+            g_ref, g_test = make_reference_and_test_tensors(
+                in_shape,
+                min=-6.0,
+                max=4.0,
+                test_dtype=dtype,
+                test_device=device,
+            )
+        else:
+            g_base_ref, g_base_test = make_reference_and_test_tensors(
+                [*in_shape[:-1], 2 * in_shape[-1]],
+                test_dtype=dtype,
+                test_device=device,
+            )
+            g_ref = g_base_ref[..., ::2]
+            g_test = g_base_test[..., ::2]
+            g_ref.retain_grad()
+            g_test.retain_grad()
+            assert not g_test.is_contiguous()
+
+        # Plain PyTorch implementation
+        var_ref = x_ref.square().mean(dim=-1, keepdim=True)
+        gamma_ref = 1 + w_ref if zero_centered_gamma else w_ref
+        n_ref = x_ref * torch.rsqrt(eps + var_ref) * gamma_ref
+        # The op rounds the normalized output and its gradient to the input
+        # dtype around the gate (as the unfused Megatron path does), so the
+        # reference does too.
+        n_ref = n_ref + (n_ref.to(dtype).to(n_ref.dtype) - n_ref).detach()
+        n_ref.register_hook(lambda grad: grad.to(dtype).to(grad.dtype))
+        y_ref = n_ref * torch.nn.functional.silu(g_ref)
+        y_ref.backward(dy_ref)
+
+        # Implementation with fusible operation
+        op = te_ops.RMSNorm(
+            weight_shape,
+            eps=eps,
+            device=device,
+            dtype=dtype,
+            zero_centered_gamma=zero_centered_gamma,
+        )
+        with torch.no_grad():
+            op.weight.copy_(w_test)
+            del w_test
+        forward = te_ops.Sequential(op, te_ops.SiLUGate())
+        y_test = forward(x_test, g_test)
+        y_test.backward(dy_test)
+
+        # Expected numerical error
+        tols = dtype_tols(dtype)
+
+        # Check results
+        y_test = y_test.to(dtype=torch.float64, device="cpu")
+        dx_test = x_test.grad.to(dtype=torch.float64, device="cpu")
+        dw_test = op.weight.grad.to(dtype=torch.float64, device="cpu")
+        dg_test = g_test.grad.to(dtype=torch.float64, device="cpu")
+        torch.testing.assert_close(y_test, y_ref, **tols)
+        torch.testing.assert_close(dx_test, x_ref.grad, **tols)
+        torch.testing.assert_close(dw_test, w_ref.grad, **tols)
+        torch.testing.assert_close(dg_test, g_ref.grad, **tols)
+
     @pytest.mark.parametrize("in_shape", ((32,), (6, 16, 64), (32, 64)))
     @pytest.mark.parametrize("dtype", _dtypes)
     def test_l2normalization(

transformer_engine/pytorch/ops/basic/__init__.py

@@ -36,6 +36,7 @@
 from .reduce_scatter import ReduceScatter
 from .reshape import Reshape
 from .rmsnorm import RMSNorm
+from .silu_gate import SiLUGate
 from .swiglu import (
     ClampedSwiGLU,
     ScaledClampedQGeGLU,

transformer_engine/pytorch/ops/basic/silu_gate.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Fusable operation for SiLU gating."""
+
+from __future__ import annotations
+from collections.abc import Iterable, Sequence
+from typing import Any, Optional
+
+import torch
+
+from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload
+from ...tensor import Quantizer
+from ..op import BasicOperation, OperationContext
+from .._common import maybe_autocast_dtype, maybe_dequantize
+
+
+class SiLUGate(BasicOperation):
+    r"""Multiply input by the SiLU activation of a gate tensor
+
+    This operation takes the main input :math:`x` and an extra gate
+    tensor :math:`g` and computes
+
+    .. math::
+        y = x * \mathrm{SiLU}(g) = \frac{x g}{1 + \exp(-g)}
+
+    The gate tensor is an extra operation input and must have the same
+    shape as the main input. Computation is performed in FP32 and the
+    output is returned in the input dtype.
+
+    This operation is typically used after a normalization operation to
+    implement gated RMSNorm, as in Mamba2 and Qwen3-Next:
+
+    .. code-block:: python
+
+        op = Sequential(RMSNorm(hidden_size), SiLUGate())
+
+    Output quantization is not supported.
+
+    """
+
+    num_extra_inputs: int = 1
+
+    def fuser_forward(
+        self,
+        basic_op_ctxs: list[OperationContext],
+        input_: torch.Tensor,
+        *,
+        basic_op_extra_inputs: list[tuple[torch.Tensor, ...]],
+        prev_op_grad_output_quantizer: Optional[Quantizer],
+        next_op_input_quantizer: Optional[Quantizer],
+        basic_op_kwargs: list[dict[str, Any]],
+    ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]:
+        # Output quantization is not supported. Return an unquantized
+        # tensor and leave quantization to the next operation.
+        del prev_op_grad_output_quantizer, next_op_input_quantizer, basic_op_kwargs
+
+        # Check input tensors
+        x = maybe_dequantize(input_.contiguous())
+        gate = maybe_dequantize(basic_op_extra_inputs[0][0].contiguous())
+        if gate.shape != x.shape:
+            raise ValueError(
+                f"Input tensor (shape={tuple(x.shape)}) "
+                f"and gate tensor (shape={tuple(gate.shape)}) are not compatible"
+            )
+        dtype = maybe_autocast_dtype(default_dtype=x.dtype)
+        x = maybe_dequantize(x, dtype)
+
+        # Compute SiLU gate in FP32
+        y = (x.float() * torch.nn.functional.silu(gate.float())).to(dtype)
+
+        # Save state for backward pass
+        ctx = basic_op_ctxs[0]
+        if ctx.requires_grad:
+            if is_cpu_offload_enabled():
+                mark_activation_offload(x, gate)
+            ctx.save_for_backward(x, gate)
+
+        return y, [()]
+
+    def fuser_backward(
+        self,
+        basic_op_ctxs: list[OperationContext],
+        grad_output: torch.Tensor,
+        *,
+        basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]],
+    ) -> tuple[
+        torch.Tensor,
+        Iterable[Iterable[Optional[torch.Tensor]]],
+        Iterable[Iterable[Optional[torch.Tensor]]],
+    ]:
+        del basic_op_grad_extra_outputs
+
+        # Saved tensors from forward pass
+        ctx = basic_op_ctxs[0]
+        x, gate = ctx.saved_tensors
+
+        dy = maybe_dequantize(grad_output.contiguous(), x.dtype)
+
+        # Gradients of the SiLU gate, computed in FP32
+        x_fp32 = x.float()
+        gate_fp32 = gate.float()
+        dy_fp32 = dy.float()
+        sigmoid = torch.sigmoid(gate_fp32)
+        silu_grad = sigmoid * (1 + gate_fp32 * (1 - sigmoid))
+        dx = (dy_fp32 * torch.nn.functional.silu(gate_fp32)).to(x.dtype)
+        dgate = (dy_fp32 * x_fp32 * silu_grad).to(gate.dtype)
+
+        return dx, [()], [(dgate,)]