[CuTeDSL] Add grouped SiTU-GLU activation

docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md

@@ -12,7 +12,7 @@ For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point
 
 **Unified Grouped GEMM + dGLU fusion**: one public class and wrapper select a
 plain BF16 or legacy block-scaled grouped GEMM fused with a dGLU backward
-epilogue (dSwiGLU or dGeGLU) on NVIDIA Blackwell GPUs (SM100+). The operation
+epilogue (dSwiGLU, dGeGLU, or block-scaled dSiTU-GLU) on NVIDIA Blackwell GPUs (SM100+). The operation
 is implemented with CUTLASS/CuTe DSL.
 
 This is a **unified API** that supports both weight layout modes:
@@ -22,6 +22,7 @@ This is a **unified API** that supports both weight layout modes:
 And both backward activation functions:
 - **dSwiGLU**: `act_func="dswiglu"` (default)
 - **dGeGLU**: `act_func="dgeglu"`
+- **dSiTU-GLU**: `act_func="dsituglu"` (block-scaled SM100/SM103 only)
 
 Groups are contiguous in the M dimension and described by `padded_offsets` (cumulative aligned end offsets).
 
@@ -34,6 +35,7 @@ Groups are contiguous in the M dimension and described by `padded_offsets` (cumu
 
 Mixed families and unsupported pairs are rejected before allocation or
 compilation. Each backend's argument contract is described below.
+`dsituglu` is not available on the BF16 or Rubin backends.
 
 ## BF16 contract
 
@@ -125,6 +127,32 @@ The block-scaled backend performs:
 
 ### Equations
 
+For dSiTU-GLU, define
+
+$$
+T_g=\beta_1\tanh(G/\beta_1)\sigma(G),\qquad
+T_u=\beta_2\tanh(U/\beta_2).
+$$
+
+The fused backward computes
+
+$$
+\frac{\partial T_g}{\partial G}=
+(1-\tanh^2(G/\beta_1))\sigma(G)+
+\beta_1\tanh(G/\beta_1)\sigma(G)(1-\sigma(G)),
+$$
+
+$$
+\frac{\partial T_u}{\partial U}=1-\tanh^2(U/\beta_2),
+$$
+
+and returns `ref * prob * T_u * dT_g/dG` and
+`ref * prob * T_g * dT_u/dU`. `dprob` accumulates the reduction of
+`ref * T_g * T_u` across the output columns in 32-column chunks, producing
+shape `(valid_m, 1, 1)`.
+The beta values are compile-time specialization values and therefore belong to
+the dGLU compiled-kernel cache key.
+
 - **Inputs**
   - `A`: contiguous activation tensor across all groups, shape `(valid_m, K, 1)`
   - `B` (dense): weight tensor across all groups, shape `(N, K, L)`
@@ -492,10 +520,14 @@ Providing both or neither raises `ValueError`.
   - `TILE_N = 256`
 - `cluster_shape_mn`: Thread Block cluster shape. Default: `(2, 1)` when `TILE_M=256`, `(1, 1)` otherwise
 - `sf_vec_size`: Scale factor vector size. `{16, 32}`. Default: `16`
-- `vector_f32`: Enable packed f32 operations. Default: `False`
+- `vector_f32`: Enable packed f32 operations for dSwiGLU and dGeGLU. Default:
+  `False`. K3-default dSiTU-GLU (`situ_beta1=4.0`) always uses its packed FP32x2
+  specialization; non-default dSiTU-GLU uses scalar FP32.
 - `m_aligned`: Must be `256`. Default: `256`
 - `discrete_col_sfd`: Generate discrete col-major scale factors. Default: `False`
-- `act_func`: Backward activation function. `"dswiglu"` (default) or `"dgeglu"`
+- `act_func`: Backward activation function. `"dswiglu"` (default), `"dgeglu"`, or block-scaled `"dsituglu"`
+- `situ_beta1`: Positive finite gate tanh scale for dSiTU-GLU. Default: `4.0`
+- `situ_beta2`: Positive finite up-branch tanh scale for dSiTU-GLU. Default: `25.0`
 - `b_major` (discrete only): B tensor major dimension. `"k"` (default) or `"n"`. Must be `"k"` for FP4.
 - `epilogue_op`: Optional post-processing. `None` (default), `"identity"`, `"relu"`, or `"srelu"`
 

docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md

@@ -12,7 +12,7 @@ For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point
 
 **Unified Grouped GEMM + GLU fusion**: one public class and wrapper select a
 plain BF16 or legacy block-scaled grouped GEMM fused with a GLU epilogue
-(SwiGLU or GeGLU) on NVIDIA Blackwell GPUs (SM100+). The operation is
+(SwiGLU, GeGLU, or block-scaled SiTU-GLU) on NVIDIA Blackwell GPUs (SM100+). The operation is
 implemented with CUTLASS/CuTe DSL.
 
 This is a **unified API** that supports both weight layout modes:
@@ -22,6 +22,7 @@ This is a **unified API** that supports both weight layout modes:
 And both activation functions:
 - **SwiGLU**: `act_func="swiglu"` (default)
 - **GeGLU**: `act_func="geglu"`
+- **SiTU-GLU**: `act_func="situglu"` (block-scaled SM100/SM103 only)
 
 Groups are contiguous in the M dimension and described by `padded_offsets` (cumulative aligned end offsets).
 
@@ -34,6 +35,7 @@ Groups are contiguous in the M dimension and described by `padded_offsets` (cumu
 
 Mixed families and unsupported pairs are rejected before allocation or
 compilation. Each backend's argument contract is described below.
+SiTU-GLU is not available on the BF16 or Rubin backends.
 
 ## BF16 contract
 
@@ -91,13 +93,26 @@ always `None`; `c_tensor` is `None` unless `generate_c=True`.
 
 The block-scaled backend performs:
 1. **Block-scaled grouped GEMM**: Low-precision GEMM (FP4, FP8) with per-block scale factors across multiple expert groups
-2. **GLU activation**: Fused SwiGLU or GeGLU activation applied to the GEMM output
+2. **GLU activation**: Fused SwiGLU, GeGLU, or SiTU-GLU activation applied to the GEMM output
 3. **Optional quantized output**: Produces row and column scale factors for downstream quantization
 
 ### Shapes
 
 ### Equations
 
+For SiTU-GLU, with gate branch `G` and up branch `U`, the fused epilogue computes
+
+$$
+D = \mathrm{prob}\,
+    \left[\beta_1\tanh(G/\beta_1)\sigma(G)\right]
+    \left[\beta_2\tanh(U/\beta_2)\right].
+$$
+
+where `beta_1 = situ_beta1` and `beta_2 = situ_beta2`, with defaults
+`beta_1 = 4.0` and `beta_2 = 25.0`. `situ_beta1` specializes the compiled kernel
+and is part of its cache key; `situ_beta2` is a runtime FP32 scalar and does not
+create a new compiled-kernel cache entry.
+
 - **Inputs**
   - `A`: contiguous activation tensor across all groups, shape `(valid_m, K, 1)`
   - `B` (dense): weight tensor across all groups, shape `(N, K, L)`
@@ -470,7 +485,9 @@ Providing both or neither raises `ValueError`.
 - `vector_f32`: Enable packed f32 operations. Default: `False`
 - `m_aligned`: Must be `256` (FIX_PAD_SIZE). Default: `256`
 - `discrete_col_sfd`: Generate discrete col-major scale factors. Default: `False`
-- `act_func`: Activation function. `"swiglu"` (default) or `"geglu"`
+- `act_func`: Activation function. `"swiglu"` (default), `"geglu"`, or block-scaled `"situglu"`
+- `situ_beta1`: Positive finite gate tanh scale for SiTU-GLU. Default: `4.0`
+- `situ_beta2`: Positive finite up-branch tanh scale for SiTU-GLU. Default: `25.0`
 - `b_major` (discrete only): B tensor major dimension. `"k"` (default) or `"n"`. Must be `"k"` for FP4.
 
 ### Wrapper-specific Parameters

docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md

@@ -8,15 +8,15 @@ JAX arrays are **not supported**: this fusion is block-scaled-only and its manda
 
 ## Overview
 
-**Grouped GEMM + GLU + Hadamard fusion**: A contiguous grouped block-scaled GEMM fused with a GLU epilogue, a 16-wide Hadamard transform, and per-expert `amax` reduction on NVIDIA Blackwell GPUs (SM100+), designed for MoE-style workloads. Groups are contiguous in the `M` dimension and described by `padded_offsets`.
+**Grouped GEMM + GLU + Hadamard fusion**: A contiguous grouped block-scaled GEMM fused with a GLU epilogue, a 16-wide Hadamard transform for post-RHT amax computation, and per-expert `amax` reductions on NVIDIA Blackwell GPUs (SM100+), designed for MoE-style workloads. Groups are contiguous in the `M` dimension and described by `padded_offsets`.
 
 This frontend integration is currently wired for the fp4 input path.
 
 This kernel performs:
 1. **Block-scaled grouped GEMM** over contiguous expert ranges
-2. **GLU epilogue** using per-row `prob`
-3. **Hadamard transform** across the post-GLU output
-4. **Per-expert amax reduction** on the final output
+2. **GLU epilogue** using per-row `prob` with SwiGLU, GeGLU, SiTU-GLU, or SReLU
+3. **Hadamard transform** over 16-token groups of the post-GLU output
+4. **Per-expert amax reductions** before and after the Hadamard transform
 
 ### Shapes
 
@@ -33,8 +33,9 @@ This kernel performs:
 
 - **Outputs**
   - `C`: intermediate GEMM result before GLU/Hadamard, shape `(valid_m, N, 1)`
-  - `D`: output after GLU and Hadamard, shape `(valid_m, N / 2, 1)`
-  - `Amax`: per-expert amax, shape `(L, 1)` when `D` is fp16/bf16
+  - `D`: activation output before the Hadamard transform, shape `(valid_m, N / 2, 1)` for GLU activations and `(valid_m, N, 1)` for SReLU
+  - `Amax`: per-expert amax of `D`, shape `(L, 1)` when `D` is fp16/bf16
+  - `PostRhtAmax`: per-expert amax after the normalized Hadamard transform, shape `(L, 1)` when `D` is fp16/bf16
 
 `L` is the expert count and `valid_m = padded_offsets[-1]`.
 
@@ -64,13 +65,23 @@ $$
 X[:, bG:(b+1)G] = \mathrm{prob} \cdot (U_b + 1) \cdot G_b \cdot \sigma(1.702 \cdot G_b)
 $$
 
-Apply the fixed Hadamard matrix `H` of size `16 x 16` blockwise over the output:
+For **SiTU-GLU** (`act_func="situglu"`):
 
 $$
-D = X \cdot H
+X[:, bG:(b+1)G] = \mathrm{prob}\,
+    \left[\beta_1\tanh(G_b/\beta_1)\sigma(G_b)\right]
+    \left[\beta_2\tanh(U_b/\beta_2)\right].
 $$
 
-When `D` is fp16/bf16, the kernel also emits per-expert `Amax`.
+Here `beta_1 = situ_beta1` and `beta_2 = situ_beta2`, with defaults `4.0` and `25.0`. `situ_beta1` specializes the compiled kernel and is part of its cache key. `situ_beta2` is a runtime FP32 scalar, so changing it does not create a new compiled-kernel cache entry.
+
+The returned `D` is `X`. For NVFP4 quantization, the kernel also applies the normalized fixed Hadamard matrix `H` of size `16 x 16` over 16-token groups within each expert and reduces its absolute maximum:
+
+$$
+\mathrm{PostRhtAmax}_g = \max \left|\mathrm{RHT}_{16}(X_g)\right|.
+$$
+
+When `D` is fp16/bf16, the kernel emits both `Amax`, computed from the untransformed `D`, and `PostRhtAmax`. The transformed values are not materialized as another output tensor; the post-RHT amax is intended for the downstream NVFP4 quantization step.
 
 ### Diagram
 
@@ -89,14 +100,11 @@ A (valid_m×K×1), SFA     B (N×K×L), SFB       padded_offsets
                      | GLU over paired 32-col blocks
                      | with per-row prob
                      v
-                 X (valid_m×N/2×1)
-                     |
-                     | blockwise Hadamard(16)
-                     v
                  D (valid_m×N/2×1)
+                     |\
+                     | +--> Amax (L×1)
                      |
-                     v
-                 Amax (L×1)
+                     +----> normalized RHT(16) --> PostRhtAmax (L×1)
 ```
 
 ---
@@ -126,11 +134,13 @@ result = grouped_gemm_glu_hadamard_wrapper_sm100(
     sf_vec_size=16,
     vector_f32=False,
     m_aligned=256,
-    act_func="swiglu",
+    act_func="situglu",
+    situ_beta1=4.0,
+    situ_beta2=25.0,
     current_stream=None,
 )
 
-c_tensor, d_tensor, amax_tensor = result
+c_tensor, d_tensor, amax_tensor, post_rht_amax_tensor = result
 ```
 
 The wrapper constructs the fixed Hadamard matrix internally.
@@ -151,14 +161,17 @@ op = GroupedGemmGluHadamardSm100(
     sample_alpha=alpha,
     sample_prob=prob,
     sample_amax=amax,
+    sample_post_rht_amax=post_rht_amax,
     sample_bias=bias,
     acc_dtype=torch.float32,
     mma_tiler_mn=(256, 256),
     cluster_shape_mn=(2, 1),
     sf_vec_size=16,
     vector_f32=False,
     m_aligned=256,
-    act_func="swiglu",
+    act_func="situglu",
+    situ_beta1=4.0,
+    situ_beta2=25.0,
 )
 assert op.check_support()
 op.compile()
@@ -173,7 +186,10 @@ op.execute(
     alpha_tensor=alpha,
     prob_tensor=prob,
     amax_tensor=amax,
+    post_rht_amax_tensor=post_rht_amax,
     bias_tensor=bias,
+    situ_beta1=4.0,
+    situ_beta2=25.0,
     current_stream=None,
 )
 ```
@@ -222,12 +238,16 @@ You may optionally pass a custom `sample_hadamard` / `hadamard_tensor`, but the
   - Layout: must be `n`-major
   - Dtype: `{float16, bfloat16}`
 - Output tensor **D**: `result["d_tensor"]` (wrapper) or `sample_d` / `d_tensor` (class)
-  - Shape: `(valid_m, N / 2, 1)`
+  - Shape: `(valid_m, N / 2, 1)` for GLU activations; `(valid_m, N, 1)` for SReLU
   - Layout: must be `n`-major
   - Dtype: `{float16, bfloat16}`
 - Output tensor **Amax**: `result["amax_tensor"]` (wrapper) or `sample_amax` / `amax_tensor` (class)
   - Shape: `(L, 1)`
   - Dtype: `float32`
+- Output tensor **PostRhtAmax**: `result["post_rht_amax_tensor"]` (wrapper) or `sample_post_rht_amax` / `post_rht_amax_tensor` (class)
+  - Shape: `(L, 1)`
+  - Dtype: `float32`
+  - Semantics: per-expert amax after normalized RHT(16), for downstream NVFP4 quantization
 
 ### Common parameters
 
@@ -244,7 +264,13 @@ You may optionally pass a custom `sample_hadamard` / `hadamard_tensor`, but the
 - `m_aligned: int`
   - Must equal the kernel fixed pad size `256`
 - `act_func: str`
-  - Allowed values: `{"swiglu", "geglu"}`
+  - Allowed values: `{"swiglu", "geglu", "situglu", "srelu"}`
+- `situ_beta1: float`
+  - Positive finite gate tanh scale for SiTU-GLU; default `4.0`
+  - Compile-time specialized and included in the wrapper cache key
+- `situ_beta2: float`
+  - Positive finite up-branch tanh scale for SiTU-GLU; default `25.0`
+  - Runtime FP32 scalar; changing it reuses the compiled beta1 specialization
 - CUDA stream (`current_stream` in class API and wrapper)
 
 ### Wrapper return values
@@ -254,8 +280,9 @@ Returns a `TupleDict` with keys:
 - `c_tensor`
 - `d_tensor`
 - `amax_tensor`
+- `post_rht_amax_tensor`
 
-Tuple unpacking order is: `(c_tensor, d_tensor, amax_tensor)`.
+Tuple unpacking order is: `(c_tensor, d_tensor, amax_tensor, post_rht_amax_tensor)`.
 
 ---