[JAX] Add SiTU-GLU support to JAX

tests/jax/test_custom_call_compute.py

@@ -199,6 +199,7 @@ def assert_dequantized_grouped_scaled_tensor(
     ("squared_relu",),
     ("squared_relu", "linear"),
     ("clamped_silu", "clamped_linear"),
+    ("situ", "situ_linear"),
 ]
 
 ACTIVATION_TYPES = {
@@ -210,6 +211,17 @@ def assert_dequantized_grouped_scaled_tensor(
 }
 
 
+def make_activation_params(activation_type):
+    """Create non-default parameters for configurable activation tests."""
+    if activation_type == ("clamped_silu", "clamped_linear"):
+        return tex.activation.ActivationParams.create(
+            activation_type, limit=0.75, alpha=1.702, glu_linear_offset=0.5
+        )
+    if activation_type == ("situ", "situ_linear"):
+        return tex.activation.ActivationParams.create(activation_type, beta1=2.0, beta2=8.0)
+    return None
+
+
 class TestActivation:
     def ref_act(self, x, activation_type, act_params):
         return _jax_act_lu(x, activation_type, act_params=act_params).data
@@ -228,6 +240,26 @@ def primitive_func(self, inputs, activation_type, quantizer, act_params):
         )
         return jnp.mean(out)
 
+    @pytest.mark.parametrize("betas", [(4.0, 25.0), (2.0, 8.0)])
+    def test_jax_situglu_reference(self, betas):
+        beta1, beta2 = betas
+        x = jnp.arange(24, dtype=jnp.float32).reshape(4, 2, 3) / 3.0 - 4.0
+        params = tex.activation.ActivationParams.create(
+            ("situ", "situ_linear"), beta1=beta1, beta2=beta2
+        )
+
+        output = _jax_act_lu(x, ("situ", "situ_linear"), act_params=params).data
+        gate, up = x[:, 0, :], x[:, 1, :]
+        expected = (beta1 * jnp.tanh(gate / beta1) * jax.nn.sigmoid(gate)) * (
+            beta2 * jnp.tanh(up / beta2)
+        )
+
+        assert_allclose(output, expected, dtype=x.dtype)
+        assert hash(params) == hash(params)
+        ffi_params = params.to_ffi_lowering_dict()["situglu"]
+        assert float(ffi_params["beta1"]) == pytest.approx(beta1)
+        assert float(ffi_params["beta2"]) == pytest.approx(beta2)
+
     @pytest_parametrize_wrapper("shape", ALL_ACTIVATION_SHAPES)
     @pytest_parametrize_wrapper(
         "activation_type",
@@ -244,16 +276,7 @@ def test_act_grad(self, shape, activation_type):
         value_n_grad_primitive_func = jit(
             value_and_grad(self.primitive_func, (0,)), static_argnums=(1, 3)
         )
-        act_args = (
-            {"limit": 0.75, "alpha": 1.702, "glu_linear_offset": 0.5}
-            if activation_type == ("clamped_silu", "clamped_linear")
-            else {}
-        )
-        act_params = (
-            tex.activation.ActivationParams.create(activation_type=activation_type, **act_args)
-            if activation_type == ("clamped_silu", "clamped_linear")
-            else None
-        )
+        act_params = make_activation_params(activation_type)
         prim_out, (prim_grad,) = value_n_grad_primitive_func(x, activation_type, None, act_params)
         ref_out, (ref_grad,) = self.value_n_grad_ref_func(x, activation_type, act_params)
         assert_allclose(prim_out, ref_out, dtype=x.dtype)
@@ -284,17 +307,7 @@ def test_act_grad_with_tensor_scaling_fp8(
             q_dtype=output_type,
             q_layout=QuantizeLayout.ROWWISE,
         )
-        act_args = (
-            {"limit": 0.75, "alpha": 1.702}
-            if activation_type == ("clamped_silu", "clamped_linear")
-            else {}
-        )
-
-        act_params = (
-            tex.activation.ActivationParams.create(activation_type=activation_type, **act_args)
-            if activation_type == ("clamped_silu", "clamped_linear")
-            else None
-        )
+        act_params = make_activation_params(activation_type)
         prim_out, (prim_grad,) = value_n_grad_primitive_func(
             x, activation_type, quantizer, act_params
         )
@@ -327,16 +340,7 @@ def test_act_forward_with_tensor_scaling_fp8(
             q_dtype=output_type,
             q_layout=q_layout,
         )
-        act_args = (
-            {"limit": 0.75, "alpha": 1.702}
-            if activation_type == ("clamped_silu", "clamped_linear")
-            else {}
-        )
-        act_params = (
-            tex.activation.ActivationParams.create(activation_type=activation_type, **act_args)
-            if activation_type == ("clamped_silu", "clamped_linear")
-            else None
-        )
+        act_params = make_activation_params(activation_type)
         te_output = tex.act_lu(x, activation_type, te_quantizer, act_params)
         jax_output = _jax_act_lu(x, activation_type, jax_quantizer, act_params)
         assert_bitwise_scaled_tensors(te_output, jax_output)
@@ -358,20 +362,18 @@ def test_act_forward_with_block_scaling_fp8(
         quantizer = QuantizerFactory.create(
             scaling_mode=ScalingMode.MXFP8_1D_SCALING, q_dtype=output_type, q_layout=q_layout
         )
-        act_args = (
-            {"limit": 0.75, "alpha": 1.702}
-            if activation_type == ("clamped_silu", "clamped_linear")
-            else {}
-        )
-        act_params = (
-            tex.activation.ActivationParams.create(activation_type=activation_type, **act_args)
-            if activation_type == ("clamped_silu", "clamped_linear")
-            else None
-        )
+        act_params = make_activation_params(activation_type)
         output = tex.act_lu(x, activation_type, quantizer, act_params)
         ref_out = self.ref_act(x, activation_type, act_params)
         assert_dequantized_scaled_tensor(output, ref_out)
 
+    @pytest.mark.parametrize("name", ["beta1", "beta2"])
+    @pytest.mark.parametrize("value", [0.0, -1.0, jnp.inf, -jnp.inf, jnp.nan])
+    def test_situglu_invalid_params(self, name, value):
+        kwargs = {name: value}
+        with pytest.raises(ValueError, match=rf"{name} must be finite and positive"):
+            tex.activation.ActivationParams.create(("situ", "situ_linear"), **kwargs)
+
 
 NORM_OUTPUT_DTYPES = {
     "L0": [jnp.float8_e4m3fn],
@@ -1631,7 +1633,9 @@ def ref_func(x, w, gamma, beta):
 
     @pytest.mark.skipif(not is_fp8_supported, reason=fp8_unsupported_reason)
     @pytest.mark.parametrize("m,n,k", [(64, 128, 128)])
-    @pytest.mark.parametrize("activation_type", [("gelu",), ("gelu", "linear")])
+    @pytest.mark.parametrize(
+        "activation_type", [("gelu",), ("gelu", "linear"), ("situ", "situ_linear")]
+    )
     @pytest_parametrize_wrapper("recipe", supported_recipes)
     @pytest.mark.parametrize("norm_type", ["layernorm", "rmsnorm"])
     @pytest_parametrize_wrapper("use_bias", [True, False])
@@ -1670,6 +1674,10 @@ def test_layernorm_mlp_grad(
                 x=QuantizeMeta(), kernel=QuantizeMeta(), grad=QuantizeMeta()
             ),
         )
+        activation_params = (
+            {"beta1": 2.0, "beta2": 8.0} if activation_type == ("situ", "situ_linear") else None
+        )
+        ref_activation_params = make_activation_params(activation_type)
 
         if norm_type == "layernorm":
             beta = jax.random.normal(subkeys[3], (k,)).astype(jnp.bfloat16)
@@ -1688,6 +1696,7 @@ def prim_func(x, gamma, kernel_1, kernel_2, bias_1, bias_2):
                     zero_centered_gamma=zero_centered_gamma,
                     epsilon=eps,
                     activation_type=activation_type,
+                    activation_params=activation_params,
                     quantizer_sets=quantizer_sets,
                 )
             )
@@ -1701,7 +1710,7 @@ def _ref_func_impl(x, gamma, kernel_1, kernel_2, bias_1, bias_2):
                 bias_1_shape = (1,) * (linear_1_out.ndim - bias_1.ndim) + bias_1.shape
                 linear_1_out += jnp.reshape(bias_1, bias_1_shape)
 
-            x = _jax_act_lu(linear_1_out, activation_type).data
+            x = _jax_act_lu(linear_1_out, activation_type, act_params=ref_activation_params).data
             linear_2_out = jax.lax.dot_general(x, kernel_2, (((1,), (0,)), ((), ())))
             if use_bias:
                 bias_2_shape = (1,) * (linear_2_out.ndim - bias_2.ndim) + bias_2.shape

tests/jax/test_distributed_layernorm_mlp.py

@@ -337,6 +337,12 @@ def _test_layernorm_mlp(
     ):
         batch, seqlen, hidden_in = input_shape
         layernorm_type = "rmsnorm"
+        if activation_type == ("situ", "situ_linear"):
+            activation_params = {"beta1": 2.0, "beta2": 8.0}
+        elif activation_type == ("clamped_silu", "clamped_linear"):
+            activation_params = {"limit": 0.75, "alpha": 1.702, "glu_linear_offset": 0.5}
+        else:
+            activation_params = None
 
         rng = jax.random.PRNGKey(0)
         subkeys = jax.random.split(rng, 3)
@@ -353,6 +359,7 @@ def _test_layernorm_mlp(
                     layernorm_type=layernorm_type,
                     intermediate_dim=INTERMEDIATE,
                     activations=activation_type,
+                    activation_params=activation_params,
                     use_bias=use_bias,
                     return_layernorm_output=True,
                 )
@@ -372,6 +379,7 @@ def _test_layernorm_mlp(
                     layernorm_type=layernorm_type,
                     intermediate_dim=INTERMEDIATE,
                     activations=activation_type,
+                    activation_params=activation_params,
                     scale_axes=LN_SCALE_AXES,
                     ln_bias_axes=LN_BIAS_AXES,
                     kernel_axes_1=KERNEL_1_AXES,
@@ -431,7 +439,15 @@ def _test_layernorm_mlp(
 
     @pytest_parametrize_wrapper("input_shape", INPUT_SHAPE)
     @pytest_parametrize_wrapper("mesh_config", generate_fsdp_and_tpsp_configs())
-    @pytest_parametrize_wrapper("activation_type", [("gelu",), ("silu", "linear")])
+    @pytest_parametrize_wrapper(
+        "activation_type",
+        [
+            ("gelu",),
+            ("silu", "linear"),
+            ("clamped_silu", "clamped_linear"),
+            ("situ", "situ_linear"),
+        ],
+    )
     @pytest_parametrize_wrapper("dtype", DTYPES)
     @pytest_parametrize_wrapper("use_bias", [True, False])
     @pytest_parametrize_wrapper("with_jax_gemm", [False, True])

transformer_engine/common/include/transformer_engine/activation.h

@@ -40,7 +40,8 @@ enum class NVTE_Activation_Type {
   QGEGLU,
   SRELU,
   SREGLU,
-  CLAMPED_SWIGLU
+  CLAMPED_SWIGLU,
+  SITUGLU
 };
 
 /*! \brief Computes the GeLU activation of the input.

transformer_engine/jax/activation.py

@@ -17,23 +17,30 @@
 from .quantize.quantizer import Quantizer
 
 
+ActivationParams = tex.activation.ActivationParams
+
+__all__ = ["activation", "ActivationParams"]
+
+
 def activation(
     x: jnp.ndarray,
     activation_type: Sequence[Union[str, Callable]],
     quantizer: Optional[Quantizer] = None,
-    act_params: Optional[tex.activation.ActivationParams] = None,
+    act_params: Optional[ActivationParams] = None,
 ) -> jnp.ndarray:
     """Apply activation functions to input tensor with optional quantization.
 
     This function applies a sequence of activation functions to the input tensor.
     It supports string-based activation types (e.g., 'relu', 'gelu', ('gelu', 'linear')).
+    SiTU-GLU is selected with ``('situ', 'situ_linear')``; ``situ_linear`` computes
+    the soft-capped up branch and is not an identity function.
 
     Args:
         x: Input tensor to apply activations to
         activation_type: Sequence of activation functions
         quantizer: Optional quantizer for quantizing the output
-        act_params: Optional activation parameters. Currently used
-        just for ClampedSwiGLU.
+        act_params: Optional parameters for configurable activations such as
+            ClampedSwiGLU and SiTU-GLU.
 
     Returns:
         Activated output tensor
@@ -54,8 +61,8 @@ def _activation(x, activation_type, quantizer, act_params):
         x: Input tensor
         activation_type: Sequence of activation functions
         quantizer: Optional quantizer
-        act_params: Optional activation parameters. Currently used
-        just for ClampedSwiGLU.
+        act_params: Optional parameters for configurable activations such as
+            ClampedSwiGLU and SiTU-GLU.
 
     Returns:
         Activated tensor
@@ -71,8 +78,8 @@ def _activation_fwd_rule(x, activation_type, quantizer, act_params):
         x: Input tensor
         activation_type: Sequence of activation functions
         quantizer: Optional quantizer
-        act_params: Optional activation parameters. Currently used
-        just for ClampedSwiGLU.
+        act_params: Optional parameters for configurable activations such as
+            ClampedSwiGLU and SiTU-GLU.
 
     Returns:
         Tuple of (output, context) for backward pass
@@ -88,8 +95,8 @@ def _activation_bwd_rule(activation_type, act_params, ctx, g):
 
     Args:
         activation_type: Sequence of activation functions
-        act_params: Optional activation parameters. Currently used
-        just for ClampedSwiGLU.
+        act_params: Optional parameters for configurable activations such as
+            ClampedSwiGLU and SiTU-GLU.
         ctx: Context from forward pass
         g: Gradient from upstream