From 23fcf9ae194b40397fdc837a4a765e261f3f8179 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 18 Aug 2026 16:46:31 +0800 Subject: [PATCH 1/9] Add grouped SiTU-GLU activation Signed-off-by: Harry Zhou --- .../gemm_fusions/grouped_gemm_dglu.md | 32 +++++++- .../gemm_fusions/grouped_gemm_glu.md | 22 +++++- .../cutedsl/grouped/dglu/_blockscaled_api.py | 32 +++++++- python/cudnn/gemm/cutedsl/grouped/dglu/api.py | 47 +++++++++-- ...moe_blockscaled_grouped_gemm_dglu_dbias.py | 64 ++++++++++++++- .../cutedsl/grouped/glu/_blockscaled_api.py | 40 +++++++++- python/cudnn/gemm/cutedsl/grouped/glu/api.py | 39 +++++++++- .../moe_blockscaled_grouped_gemm_glu_bias.py | 38 ++++++++- ...est_discrete_grouped_gemm_dswiglu_utils.py | 30 +++++-- ...test_discrete_grouped_gemm_swiglu_utils.py | 14 +++- .../grouped_gemm/test_grouped_gemm_dglu.py | 78 ++++++++++++++++++- .../grouped_gemm/test_grouped_gemm_glu.py | 77 +++++++++++++++++- 12 files changed, 473 insertions(+), 40 deletions(-) diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md index b8a85fd2a..d3db216a9 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md +++ b/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,30 @@ 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 `R * prob * T_u * dT_g/dG` and +`R * prob * T_g * dT_u/dU`. `dprob` accumulates `R * T_g * T_u`. +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)` @@ -495,7 +521,9 @@ Providing both or neither raises `ValueError`. - `vector_f32`: Enable packed f32 operations. Default: `False` - `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"` diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md index eacd4a3e5..1456193bf 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md +++ b/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,25 @@ 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`, up branch `U`, `situ_beta1=4`, and +`situ_beta2=25`, 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]. +$$ + +The beta values are runtime FP32 scalars and do not create new compiled-kernel +cache entries. + - **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 +484,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 diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py index 69e01351c..55790da3f 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py @@ -5,7 +5,7 @@ This module provides a single API class that supports both contiguous (dense) and discrete weight modes for block-scaled grouped GEMM with dGLU activation -gradient (dSwiGLU / dGeGLU) in MoE (Mixture of Experts) workloads. +gradient (dSwiGLU / dGeGLU / dSiTU-GLU) in MoE (Mixture of Experts) workloads. Dense mode All expert weights are packed contiguously in a 3-D tensor (N, K, L). @@ -23,6 +23,7 @@ from ..moe_utils import MoEWeightMode from ..backend_utils import rubin_single_group_offsets_kwarg from cuda.bindings import driver as cuda +import math import os from typing import Literal, Tuple, Optional @@ -141,6 +142,8 @@ def __init__( geglu_alpha: float = 1.702, glu_clamp_max: float = 7.0, glu_clamp_min: float = -7.0, + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, ): """Initialize the GroupedGemmDgluSm100 API. @@ -177,7 +180,7 @@ def __init__( :param vector_f32: Use vectorized f32 operations :param m_aligned: Alignment for group M dimension :param discrete_col_sfd: Generate discrete col-major scale factor tensor - :param act_func: Activation function, one of "dswiglu" or "dgeglu" + :param act_func: Activation function, one of "dswiglu", "dgeglu", or "dsituglu" :param b_major: Major dimension for B tensor, one of "k" or "n" :param epilogue_op: Optional epilogue operation. Valid: None, "none", "identity", "relu", "srelu" :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing @@ -190,6 +193,8 @@ def __init__( ``act_func == "dswiglu"``. :param glu_clamp_min: Compile-time dGeGLU lower clamp. Ignored when ``act_func == "dswiglu"``. + :param situ_beta1: Compile-time gate tanh scale for dSiTU-GLU. + :param situ_beta2: Compile-time up-branch tanh scale for dSiTU-GLU. """ from cudnn.tensor_adapter import detect_framework @@ -297,6 +302,17 @@ def __init__( self.geglu_alpha = geglu_alpha self.glu_clamp_max = glu_clamp_max self.glu_clamp_min = glu_clamp_min + self.situ_beta1 = float(situ_beta1) + self.situ_beta2 = float(situ_beta2) + if self.act_func == "dsituglu": + self._value_error_if( + not math.isfinite(self.situ_beta1) or self.situ_beta1 <= 0.0, + f"situ_beta1 must be finite and positive, got {self.situ_beta1}", + ) + self._value_error_if( + not math.isfinite(self.situ_beta2) or self.situ_beta2 <= 0.0, + f"situ_beta2 must be finite and positive, got {self.situ_beta2}", + ) _reject_unsupported_rubin_glu_tune_params( self._is_rubin_kernel, self.geglu_alpha, @@ -602,8 +618,12 @@ def check_support(self) -> bool: # ---- Activation function validation ---- self._value_error_if( - self.act_func not in ["dswiglu", "dgeglu"], - f"act_func must be 'dswiglu' or 'dgeglu', got {self.act_func}", + self.act_func not in ["dswiglu", "dgeglu", "dsituglu"], + f"act_func must be 'dswiglu', 'dgeglu', or 'dsituglu', got {self.act_func}", + ) + self._not_implemented_error_if( + self._is_rubin_kernel and self.act_func == "dsituglu", + "Rubin grouped GEMM dGLU does not support dsituglu", ) # ---- Discrete-mode-specific validation ---- @@ -997,6 +1017,8 @@ def _compile_dense(self, gemm_dglu, max_active_clusters, fake_stream) -> None: "geglu_alpha": self.geglu_alpha, "glu_clamp_max": self.glu_clamp_max, "glu_clamp_min": self.glu_clamp_min, + "situ_beta1": self.situ_beta1, + "situ_beta2": self.situ_beta2, } ) _compiled_kernel = cute.compile(gemm_dglu, **compile_kwargs) @@ -1196,6 +1218,8 @@ def _compile_discrete(self, gemm_dglu, max_active_clusters, fake_stream) -> None "geglu_alpha": self.geglu_alpha, "glu_clamp_max": self.glu_clamp_max, "glu_clamp_min": self.glu_clamp_min, + "situ_beta1": self.situ_beta1, + "situ_beta2": self.situ_beta2, } ) _compiled_kernel = cute.compile(gemm_dglu, **compile_kwargs) diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/api.py b/python/cudnn/gemm/cutedsl/grouped/dglu/api.py index 4a4d80736..405386c63 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/api.py @@ -5,7 +5,7 @@ This module provides a single API class that supports both contiguous (dense) and discrete weight modes for block-scaled grouped GEMM with dGLU activation -gradient (dSwiGLU / dGeGLU) in MoE (Mixture of Experts) workloads. +gradient (dSwiGLU / dGeGLU / dSiTU-GLU) in MoE (Mixture of Experts) workloads. Dense mode All expert weights are packed contiguously in a 3-D tensor (N, K, L). @@ -20,6 +20,7 @@ from __future__ import annotations from dataclasses import dataclass, replace +import math from ..backend_utils import ( GroupedGemmBackend, @@ -107,7 +108,7 @@ class DgluCall: mma_tiler_mn: Tuple[int, int] = (256, 256) cluster_shape_mn: Optional[Tuple[int, int]] = None sf_vec_size: int = 16 - sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None vector_f32: bool = False m_aligned: int = 256 discrete_col_sfd: bool = False @@ -116,6 +117,8 @@ class DgluCall: geglu_alpha: float = 1.702 glu_clamp_max: float = 7.0 glu_clamp_min: float = -7.0 + situ_beta1: float = 4.0 + situ_beta2: float = 25.0 epilogue_op: Optional[str] = None use_dynamic_sched: bool = False use_single_group_runtime_offsets: bool = False @@ -203,6 +206,8 @@ def __init__( geglu_alpha: float = 1.702, glu_clamp_max: float = 7.0, glu_clamp_min: float = -7.0, + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, ) -> None: super().__init__() self._pending_init_kwargs = dict(locals()) @@ -483,7 +488,7 @@ def _grouped_gemm_dglu_block_scaled_call(call: DgluCall) -> TupleDict: vector_f32: Use vectorized f32 m_aligned: M alignment (must be 256) discrete_col_sfd: Generate discrete col-major scale factor tensor - act_func: Activation function ("dswiglu" or "dgeglu") + act_func: Activation function ("dswiglu", "dgeglu", or block-scaled "dsituglu") linear_offset: Linear offset matching the forward GeGLU activation, i.e. the same value used by ``grouped_gemm_glu_wrapper_sm100`` so the backward gradients are mathematically consistent. Affects @@ -500,6 +505,10 @@ def _grouped_gemm_dglu_block_scaled_call(call: DgluCall) -> TupleDict: glu_clamp_min: Lower clamp limit applied to ``up`` only in the forward GeGLU; the same limit drives the gradient mask here. Default ``-7.0``. Ignored when ``act_func == "dswiglu"``. + situ_beta1: Positive finite gate tanh scale for dSiTU-GLU. Default + ``4.0``. Compile-time specialization and part of the cache key. + situ_beta2: Positive finite up-branch tanh scale for dSiTU-GLU. Default + ``25.0``. Compile-time specialization and part of the cache key. epilogue_op: Optional epilogue operation. Valid: None, "none", "identity", "relu", "srelu" use_dynamic_sched: Enable dynamic tile scheduling for load balancing current_stream: CUDA stream @@ -546,6 +555,8 @@ def _grouped_gemm_dglu_block_scaled_call(call: DgluCall) -> TupleDict: geglu_alpha = call.geglu_alpha glu_clamp_max = call.glu_clamp_max glu_clamp_min = call.glu_clamp_min + situ_beta1 = call.situ_beta1 + situ_beta2 = call.situ_beta2 epilogue_op = call.epilogue_op use_dynamic_sched = call.use_dynamic_sched use_single_group_runtime_offsets = call.use_single_group_runtime_offsets @@ -562,14 +573,16 @@ def _grouped_gemm_dglu_block_scaled_call(call: DgluCall) -> TupleDict: glu_clamp_max, glu_clamp_min, ) - dgeglu_cache_signature = None + activation_cache_signature = None if act_func == "dgeglu": - dgeglu_cache_signature = ( + activation_cache_signature = ( float(linear_offset), float(geglu_alpha), float(glu_clamp_max), float(glu_clamp_min), ) + elif act_func == "dsituglu": + activation_cache_signature = (float(situ_beta1), float(situ_beta2)) # ---- Auto-detect weight mode ---- is_dense = b_tensor is not None @@ -676,7 +689,7 @@ def dynamic_m_tensor_signature( device_type, weight_mode, act_func, - dgeglu_cache_signature, + activation_cache_signature, epilogue_op, use_full_dynamic, a_tensor.shape[1:] if not use_full_dynamic else None, @@ -723,7 +736,7 @@ def dynamic_m_tensor_signature( device_type, weight_mode, act_func, - dgeglu_cache_signature, + activation_cache_signature, epilogue_op, *dynamic_m_tensor_signature(a_tensor, tuple(a_tensor.shape[1:]), dynamic_stride_dims=(2,)), b_shape, @@ -802,6 +815,8 @@ def dynamic_m_tensor_signature( geglu_alpha=geglu_alpha, glu_clamp_max=glu_clamp_max, glu_clamp_min=glu_clamp_min, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, ) else: api = GroupedGemmDgluSm100( @@ -840,6 +855,8 @@ def dynamic_m_tensor_signature( geglu_alpha=geglu_alpha, glu_clamp_max=glu_clamp_max, glu_clamp_min=glu_clamp_min, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, ) if not api.check_support(): @@ -914,6 +931,16 @@ def _normalize_dglu_call( b_dtype=_convert_to_cutlass_data_type(call.b_dtype) if call.b_dtype is not None else None, ) + if call.act_func not in ("dswiglu", "dgeglu", "dsituglu"): + raise ValueError(f"act_func must be 'dswiglu', 'dgeglu', or 'dsituglu', got {call.act_func}") + if call.act_func == "dsituglu": + if not math.isfinite(call.situ_beta1) or call.situ_beta1 <= 0.0: + raise ValueError(f"situ_beta1 must be finite and positive, got {call.situ_beta1}") + if not math.isfinite(call.situ_beta2) or call.situ_beta2 <= 0.0: + raise ValueError(f"situ_beta2 must be finite and positive, got {call.situ_beta2}") + if get_device_type() == "rubin": + raise NotImplementedError("Rubin grouped GEMM dGLU does not support dsituglu") + is_dense = call.b_tensor is not None is_discrete = call.b_ptrs is not None if is_dense and is_discrete: @@ -998,7 +1025,7 @@ def _normalize_dglu_call( if call.cd_major != "n": raise ValueError(f"cd_major must be 'n', got {call.cd_major}") if call.act_func not in ("dswiglu", "dgeglu"): - raise ValueError(f"act_func must be 'dswiglu' or 'dgeglu', got {call.act_func}") + raise ValueError(f"BF16 act_func must be 'dswiglu' or 'dgeglu'; dsituglu is block-scaled only, got {call.act_func}") if call.d_dtype not in (cutlass.BFloat16, cutlass.Float16, cutlass.Float32): raise ValueError(f"d_dtype must be BF16, FP16, or FP32, got {call.d_dtype}") if call.m_aligned != 256: @@ -1236,6 +1263,8 @@ def grouped_gemm_dglu_wrapper_sm100( geglu_alpha: float = 1.702, glu_clamp_max: float = 7.0, glu_clamp_min: float = -7.0, + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, epilogue_op: Optional[str] = None, use_dynamic_sched: bool = False, use_single_group_runtime_offsets: bool = False, @@ -1283,6 +1312,8 @@ def grouped_gemm_dglu_wrapper_sm100( geglu_alpha=geglu_alpha, glu_clamp_max=glu_clamp_max, glu_clamp_min=glu_clamp_min, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, epilogue_op=epilogue_op, use_dynamic_sched=use_dynamic_sched, use_single_group_runtime_offsets=use_single_group_runtime_offsets, diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py index 651f6e190..56ef5757a 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py @@ -242,7 +242,7 @@ def can_implement( result = False if not (a_major == "k"): result = False - if act_func not in ["dswiglu", "dgeglu"]: + if act_func not in ["dswiglu", "dgeglu", "dsituglu"]: result = False return result @@ -728,6 +728,8 @@ def __call__( geglu_alpha: cutlass.Constexpr = 1.702, glu_clamp_max: cutlass.Constexpr = 7.0, glu_clamp_min: cutlass.Constexpr = -7.0, + situ_beta1: cutlass.Constexpr = 4.0, + situ_beta2: cutlass.Constexpr = 25.0, ): """Execute the GEMM. @@ -742,7 +744,9 @@ def __call__( out = (clamp(up, min=glu_clamp_min, max=glu_clamp_max) + linear_offset) * silu(geglu_alpha * clamp(gate, max=glu_clamp_max)) and the backward consumes the same values plus the corresponding - clamp masks. They are ignored when ``act_func == "dswiglu"``. + clamp masks. ``situ_beta1`` and ``situ_beta2`` configure the SiTU-GLU + derivative. GeGLU parameters are ignored unless ``act_func == "dgeglu"`` + and SiTU parameters are ignored unless ``act_func == "dsituglu"``. """ # Setup static attributes before smem/grid/tma computation self.a_dtype: Type[cutlass.Numeric] = a.element_type @@ -1100,6 +1104,8 @@ class SharedStorage: geglu_alpha, glu_clamp_max, glu_clamp_min, + situ_beta1, + situ_beta2, dbias_tensor, workspace_ptr, self.cluster_layout_vmnk, @@ -1876,6 +1882,46 @@ def dswiglu( d2_vec = acc_vec_prob * swish return d1_vec, d2_vec, dprob_swiglu + @cute.jit + def dsituglu( + self, + acc_vec: cute.Tensor, + gate_vec: cute.Tensor, + up_vec: cute.Tensor, + mProb: cute.Tensor, + beta_val: Float32, + square_alpha: Float32, + beta1: Float32, + beta2: Float32, + dprob_swiglu: Optional[cute.Tensor] = None, + ): + dgate_vec = cute.make_rmem_tensor(acc_vec.shape, cutlass.Float32) + dup_vec = cute.make_rmem_tensor(acc_vec.shape, cutlass.Float32) + + for i in cutlass.range_constexpr(cute.size(acc_vec)): + grad = acc_vec[i] * square_alpha + gate = gate_vec[i].to(cutlass.Float32) * beta_val + up = up_vec[i].to(cutlass.Float32) * beta_val + gate_tanh = cute.math.tanh(gate / beta1, fastmath=True) + up_tanh = cute.math.tanh(up / beta2, fastmath=True) + sigmoid = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True)) + gate_value = beta1 * gate_tanh * sigmoid + up_value = beta2 * up_tanh + gate_grad = (cutlass.Float32(1.0) - gate_tanh * gate_tanh) * sigmoid + gate_grad = gate_grad + beta1 * gate_tanh * sigmoid * (cutlass.Float32(1.0) - sigmoid) + up_grad = cutlass.Float32(1.0) - up_tanh * up_tanh + activation_grad = grad + if cutlass.const_expr(self.has_prob): + activation_grad = grad * mProb + dgate_vec[i] = activation_grad * up_value * gate_grad + dup_vec[i] = activation_grad * gate_value * up_grad + if cutlass.const_expr(self.generate_dprob): + dprob_swiglu[i] = grad * gate_value * up_value + + if cutlass.const_expr(self.generate_dprob): + dprob_swiglu = dprob_swiglu.load() + return dgate_vec.load(), dup_vec.load(), dprob_swiglu + @cute.jit def dgeglu( self, @@ -2069,6 +2115,8 @@ def kernel( geglu_alpha: cutlass.Constexpr, glu_clamp_max: cutlass.Constexpr, glu_clamp_min: cutlass.Constexpr, + situ_beta1: cutlass.Constexpr, + situ_beta2: cutlass.Constexpr, mDbias_tensor: Optional[cute.Tensor], workspace_ptr, cluster_layout_vmnk: cute.Layout, @@ -3040,6 +3088,18 @@ def kernel( cutlass.Float32(glu_clamp_min), dprob_swiglu, ) + elif cutlass.const_expr(self.act_func == "dsituglu"): + d1_vec, d2_vec, dprob_swiglu = self.dsituglu( + acc_vec, + ab1_vec_load, + ab2_vec_load, + mProb, + beta_val, + square_alpha, + cutlass.Float32(situ_beta1), + cutlass.Float32(situ_beta2), + dprob_swiglu, + ) if cutlass.const_expr(self.generate_dprob): # dprob sum reduction diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py index 24576990d..d55454ca9 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py @@ -23,6 +23,7 @@ from ..backend_utils import rubin_single_group_offsets_kwarg from ..moe_utils import MoEWeightMode from cuda.bindings import driver as cuda +import math import os from typing import Literal, Tuple, Optional @@ -165,7 +166,7 @@ def __init__( :param vector_f32: Use vectorized f32 operations :param m_aligned: Alignment for group M dimension :param discrete_col_sfd: Generate discrete col-major scale factor tensor - :param act_func: Activation function, one of "swiglu" or "geglu" + :param act_func: Activation function, one of "swiglu", "geglu", or "situglu" :param b_major: Major dimension for B tensor, one of "k" or "n" :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing """ @@ -526,8 +527,12 @@ def check_support(self) -> bool: # ---- Activation function validation (both modes) ---- self._value_error_if( - self.act_func not in ["swiglu", "geglu"], - f"act_func must be 'swiglu' or 'geglu', got {self.act_func}", + self.act_func not in ["swiglu", "geglu", "situglu"], + f"act_func must be 'swiglu', 'geglu', or 'situglu', got {self.act_func}", + ) + self._not_implemented_error_if( + self._is_rubin_kernel and self.act_func == "situglu", + "Rubin grouped GEMM GLU does not support situglu", ) # ---- Discrete-mode-specific validation ---- @@ -915,6 +920,8 @@ def _compile_dense(self, gemm_glu, max_active_clusters, fake_stream) -> None: "geglu_alpha": cutlass.Float32(_GEGGLU_ALPHA_DEFAULT), "glu_clamp_max": cutlass.Float32(_GLU_CLAMP_MAX_DEFAULT), "glu_clamp_min": cutlass.Float32(_GLU_CLAMP_MIN_DEFAULT), + "situ_beta1": cutlass.Float32(4.0), + "situ_beta2": cutlass.Float32(25.0), } ) _compiled_kernel = cute.compile(gemm_glu, **compile_kwargs) @@ -943,6 +950,8 @@ def tensor_api( geglu_alpha: float = 1.702, glu_clamp_max: float = 7.0, glu_clamp_min: float = -7.0, + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, ) -> None: norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") kernel_args = ( @@ -976,6 +985,8 @@ def tensor_api( cutlass.Float32(geglu_alpha), cutlass.Float32(glu_clamp_max), cutlass.Float32(glu_clamp_min), + cutlass.Float32(situ_beta1), + cutlass.Float32(situ_beta2), ) self._compiled_kernel = tensor_api @@ -1114,6 +1125,8 @@ def _compile_discrete(self, gemm_glu, max_active_clusters, fake_stream) -> None: cutlass.Float32(_GEGGLU_ALPHA_DEFAULT), cutlass.Float32(_GLU_CLAMP_MAX_DEFAULT), cutlass.Float32(_GLU_CLAMP_MIN_DEFAULT), + cutlass.Float32(4.0), + cutlass.Float32(25.0), options="--enable-tvm-ffi", ) @@ -1148,6 +1161,8 @@ def tensor_api( geglu_alpha: float = 1.702, glu_clamp_max: float = 7.0, glu_clamp_min: float = -7.0, + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, ) -> None: norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") b_ptrs_addr = int(b_ptrs_device.data_ptr()) @@ -1184,6 +1199,8 @@ def tensor_api( cutlass.Float32(geglu_alpha), cutlass.Float32(glu_clamp_max), cutlass.Float32(glu_clamp_min), + cutlass.Float32(situ_beta1), + cutlass.Float32(situ_beta2), ) self._compiled_kernel = tensor_api @@ -1218,6 +1235,8 @@ def execute( geglu_alpha: float = 1.702, glu_clamp_max: float = 7.0, glu_clamp_min: float = -7.0, + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, current_stream: Optional[cuda.CUstream] = None, ) -> None: """Execute the compiled kernel. @@ -1261,6 +1280,8 @@ def execute( :param glu_clamp_min: Lower clamp limit applied only to ``up`` (the kernel never lower-clamps ``gate``). Default ``-7.0``. Ignored when ``act_func == "swiglu"``. + :param situ_beta1: Gate tanh scale for SiTU-GLU. Default ``4.0``. + :param situ_beta2: Up-branch tanh scale for SiTU-GLU. Default ``25.0``. :param current_stream: CUDA stream """ self._logger.debug("Entering execute") @@ -1279,6 +1300,15 @@ def execute( # that pre-date the explicit linear_offset kwarg. if linear_offset is None: linear_offset = 1.0 if self.act_func == "geglu" else 0.0 + if self.act_func == "situglu": + self._value_error_if( + not math.isfinite(situ_beta1) or situ_beta1 <= 0.0, + f"situ_beta1 must be finite and positive, got {situ_beta1}", + ) + self._value_error_if( + not math.isfinite(situ_beta2) or situ_beta2 <= 0.0, + f"situ_beta2 must be finite and positive, got {situ_beta2}", + ) _reject_unsupported_rubin_glu_tune_params( self._is_rubin_kernel, geglu_alpha, @@ -1315,6 +1345,8 @@ def execute( geglu_alpha=geglu_alpha, glu_clamp_max=glu_clamp_max, glu_clamp_min=glu_clamp_min, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, ) else: self._compiled_kernel( @@ -1338,6 +1370,8 @@ def execute( geglu_alpha=geglu_alpha, glu_clamp_max=glu_clamp_max, glu_clamp_min=glu_clamp_min, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, ) self._logger.debug("Execute completed") diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/api.py b/python/cudnn/gemm/cutedsl/grouped/glu/api.py index 4e0a5e653..9df4063d0 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/api.py @@ -5,7 +5,7 @@ This module provides a single API class that supports both contiguous (dense) and discrete weight modes for block-scaled grouped GEMM with GLU activation -(SwiGLU / GeGLU) in MoE (Mixture of Experts) workloads. +(SwiGLU / GeGLU / SiTU-GLU) in MoE (Mixture of Experts) workloads. Dense mode All expert weights are packed contiguously in a 3-D tensor (N, K, L). @@ -20,6 +20,7 @@ from __future__ import annotations from dataclasses import dataclass, replace +import math from ..backend_utils import ( GroupedGemmBackend, @@ -107,7 +108,7 @@ class GluCall: mma_tiler_mn: Tuple[int, int] = (256, 256) cluster_shape_mn: Optional[Tuple[int, int]] = None sf_vec_size: int = 16 - sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None vector_f32: bool = False m_aligned: int = 256 discrete_col_sfd: bool = False @@ -116,6 +117,8 @@ class GluCall: geglu_alpha: float = 1.702 glu_clamp_max: float = 7.0 glu_clamp_min: float = -7.0 + situ_beta1: float = 4.0 + situ_beta2: float = 25.0 use_dynamic_sched: bool = False use_single_group_runtime_offsets: bool = False current_stream: Optional[cuda.CUstream] = None @@ -345,6 +348,8 @@ def execute( geglu_alpha: float = 1.702, glu_clamp_max: float = 7.0, glu_clamp_min: float = -7.0, + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, current_stream: Optional[cuda.CUstream] = None, ) -> None: if self._implementation is None: @@ -410,6 +415,8 @@ def execute( geglu_alpha=geglu_alpha, glu_clamp_max=glu_clamp_max, glu_clamp_min=glu_clamp_min, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, current_stream=current_stream, ) self._is_supported = self._implementation._is_supported @@ -469,7 +476,7 @@ def _grouped_gemm_glu_block_scaled_call(call: GluCall) -> TupleDict: vector_f32: Use vectorized f32 m_aligned: M alignment (must be 256) discrete_col_sfd: Generate discrete col-major scale factor tensor - act_func: Activation function ("swiglu" or "geglu") + act_func: Activation function ("swiglu", "geglu", or block-scaled "situglu") linear_offset: Linear offset applied to the up branch in the ``act_func == "geglu"`` activation, i.e. ``out = (up + linear_offset) * silu(geglu_alpha * gate)``. Ignored @@ -493,6 +500,10 @@ def _grouped_gemm_glu_block_scaled_call(call: GluCall) -> TupleDict: kernel never lower-clamps the gate). Default ``-7.0``. Runtime parameter, intentionally not part of the cache key. Ignored when ``act_func == "swiglu"``. + situ_beta1: Positive finite gate tanh scale for SiTU-GLU. Default + ``4.0``. Runtime parameter, intentionally not part of the cache key. + situ_beta2: Positive finite up-branch tanh scale for SiTU-GLU. Default + ``25.0``. Runtime parameter, intentionally not part of the cache key. use_dynamic_sched: Enable dynamic tile scheduling for load balancing current_stream: CUDA stream @@ -536,6 +547,8 @@ def _grouped_gemm_glu_block_scaled_call(call: GluCall) -> TupleDict: geglu_alpha = call.geglu_alpha glu_clamp_max = call.glu_clamp_max glu_clamp_min = call.glu_clamp_min + situ_beta1 = call.situ_beta1 + situ_beta2 = call.situ_beta2 use_dynamic_sched = call.use_dynamic_sched use_single_group_runtime_offsets = call.use_single_group_runtime_offsets current_stream = call.current_stream @@ -833,6 +846,8 @@ def dynamic_m_tensor_signature( geglu_alpha=geglu_alpha, glu_clamp_max=glu_clamp_max, glu_clamp_min=glu_clamp_min, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, current_stream=current_stream, ) else: @@ -856,6 +871,8 @@ def dynamic_m_tensor_signature( geglu_alpha=geglu_alpha, glu_clamp_max=glu_clamp_max, glu_clamp_min=glu_clamp_min, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, current_stream=current_stream, ) @@ -880,6 +897,16 @@ def _normalize_glu_call(call: GluCall) -> tuple[GluCall, GroupedGemmBackend]: b_dtype=_convert_to_cutlass_data_type(call.b_dtype) if call.b_dtype is not None else None, ) + if call.act_func not in ("swiglu", "geglu", "situglu"): + raise ValueError(f"act_func must be 'swiglu', 'geglu', or 'situglu', got {call.act_func}") + if call.act_func == "situglu": + if not math.isfinite(call.situ_beta1) or call.situ_beta1 <= 0.0: + raise ValueError(f"situ_beta1 must be finite and positive, got {call.situ_beta1}") + if not math.isfinite(call.situ_beta2) or call.situ_beta2 <= 0.0: + raise ValueError(f"situ_beta2 must be finite and positive, got {call.situ_beta2}") + if get_device_type() == "rubin": + raise NotImplementedError("Rubin grouped GEMM GLU does not support situglu") + is_dense = call.b_tensor is not None is_discrete = call.b_ptrs is not None if is_dense and is_discrete: @@ -964,7 +991,7 @@ def _normalize_glu_call(call: GluCall) -> tuple[GluCall, GroupedGemmBackend]: if call.cd_major != "n": raise ValueError(f"cd_major must be 'n', got {call.cd_major}") if call.act_func not in ("swiglu", "geglu"): - raise ValueError(f"act_func must be 'swiglu' or 'geglu', got {call.act_func}") + raise ValueError(f"BF16 act_func must be 'swiglu' or 'geglu'; situglu is block-scaled only, got {call.act_func}") if normalized.c_dtype not in (cutlass.BFloat16, cutlass.Float16, cutlass.Float32): raise ValueError(f"c_dtype must be BF16, FP16, or FP32, got {normalized.c_dtype}") if normalized.d_dtype not in (cutlass.BFloat16, cutlass.Float16, cutlass.Float32): @@ -1186,6 +1213,8 @@ def grouped_gemm_glu_wrapper_sm100( geglu_alpha: float = 1.702, glu_clamp_max: float = 7.0, glu_clamp_min: float = -7.0, + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, use_dynamic_sched: bool = False, use_single_group_runtime_offsets: bool = False, current_stream: Optional[cuda.CUstream] = None, @@ -1242,6 +1271,8 @@ def grouped_gemm_glu_wrapper_sm100( geglu_alpha=geglu_alpha, glu_clamp_max=glu_clamp_max, glu_clamp_min=glu_clamp_min, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, use_dynamic_sched=use_dynamic_sched, use_single_group_runtime_offsets=use_single_group_runtime_offsets, current_stream=current_stream, diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py b/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py index 80b16873e..ae03e6c4f 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py @@ -381,7 +381,7 @@ def __init__( self.num_epilog_warps = len(self.epilog_warp_id) self.act_func = act_func - if act_func not in ["swiglu", "geglu"]: + if act_func not in ["swiglu", "geglu", "situglu"]: raise ValueError(f"Invalid activation function: {act_func}") def _setup_attributes(self): @@ -728,6 +728,8 @@ def __call__( geglu_alpha: cutlass.Float32 = 1.702, glu_clamp_max: cutlass.Float32 = 7.0, glu_clamp_min: cutlass.Float32 = -7.0, + situ_beta1: cutlass.Float32 = 4.0, + situ_beta2: cutlass.Float32 = 25.0, ): """Execute the GEMM. @@ -743,7 +745,13 @@ def __call__( out = (clamp(up, min=glu_clamp_min, max=glu_clamp_max) + linear_offset) * silu(geglu_alpha * clamp(gate, max=glu_clamp_max)) - They are ignored when ``act_func == "swiglu"``. + ``situ_beta1`` and ``situ_beta2`` configure SiTU-GLU: + + out = beta1 * tanh(gate / beta1) * sigmoid(gate) + * beta2 * tanh(up / beta2) + + GeGLU parameters are ignored unless ``act_func == "geglu"`` and SiTU + parameters are ignored unless ``act_func == "situglu"``. """ self.a_dtype: Type[cutlass.Numeric] = a.element_type self.b_dtype: Type[cutlass.Numeric] = a.element_type @@ -1090,6 +1098,8 @@ class SharedStorage: geglu_alpha, glu_clamp_max, glu_clamp_min, + situ_beta1, + situ_beta2, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], @@ -1553,6 +1563,26 @@ def swiglu_act(self, tCompute: cute.Tensor, acc_vec_up: cute.Tensor, acc_vec_gat if cutlass.const_expr(self.has_prob): tCompute[i] = tCompute[i] * mProb + @cute.jit + def situglu_act( + self, + tCompute: cute.Tensor, + acc_vec_up: cute.Tensor, + acc_vec_gate: cute.Tensor, + mProb: cute.Tensor, + beta1: cutlass.Float32, + beta2: cutlass.Float32, + ): + for i in cutlass.range_constexpr(cute.size(tCompute)): + gate = acc_vec_gate[i] + up = acc_vec_up[i] + gate_tanh = cute.math.tanh(gate / beta1, fastmath=True) + up_tanh = cute.math.tanh(up / beta2, fastmath=True) + sigmoid = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True)) + tCompute[i] = beta1 * gate_tanh * sigmoid * beta2 * up_tanh + if cutlass.const_expr(self.has_prob): + tCompute[i] = tCompute[i] * mProb + # GPU device kernel @cute.kernel def kernel( @@ -1598,6 +1628,8 @@ def kernel( geglu_alpha: cutlass.Float32 = 1.702, glu_clamp_max: cutlass.Float32 = 7.0, glu_clamp_min: cutlass.Float32 = -7.0, + situ_beta1: cutlass.Float32 = 4.0, + situ_beta2: cutlass.Float32 = 25.0, ): """ GPU device kernel performing the Persistent batched GEMM computation. @@ -2733,6 +2765,8 @@ def kernel( self.geglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb, linear_offset, geglu_alpha) elif cutlass.const_expr(self.act_func == "swiglu"): self.swiglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb) + elif cutlass.const_expr(self.act_func == "situglu"): + self.situglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb, situ_beta1, situ_beta2) # # Generate amax diff --git a/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_utils.py b/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_utils.py index 6a01995b5..4da4df656 100644 --- a/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_utils.py +++ b/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_utils.py @@ -300,6 +300,9 @@ def run_discrete_dswiglu_ref( d_dtype: torch.dtype = torch.float32, sf_vec_size: int = 16, sf_dtype: torch.dtype = torch.float8_e8m0fnu, + act_func: str = "dswiglu", + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, ) -> Dict[str, torch.Tensor]: num_experts = len(b_ref_list) n = b_ref_list[0].shape[0] @@ -336,16 +339,28 @@ def run_discrete_dswiglu_ref( c_input = c_full.index_select(dim=1, index=dest_idx_ab) c_gate = c_full.index_select(dim=1, index=dest_idx_glu) sig = torch.sigmoid(c_gate) - swish = c_gate * sig + if act_func == "dsituglu": + gate_tanh = torch.tanh(c_gate / situ_beta1) + up_tanh = torch.tanh(c_input / situ_beta2) + gate_value = situ_beta1 * gate_tanh * sig + up_value = situ_beta2 * up_tanh + gate_grad = (1.0 - gate_tanh.square()) * sig + gate_grad = gate_grad + situ_beta1 * gate_tanh * sig * (1.0 - sig) + up_grad = 1.0 - up_tanh.square() + ref_dprob = gate_value * up_value * ref + prob = prob_tensor.expand(-1, n, -1) + ab = ref * prob * gate_value * up_grad + dswiglu = ref * prob * up_value * gate_grad + else: + swish = c_gate * sig + ref_dprob = swish * c_input * ref + prob = prob_tensor.expand(-1, n, -1) + ab = ref * prob * swish + dswiglu = ref * prob * c_input * sig * (1 + c_gate * (1 - sig)) - ref_dprob = swish * c_input * ref chunk_sums = [torch.sum(chunk, dim=1, keepdim=True) for chunk in torch.split(ref_dprob, 32, dim=1)] ref_tensors["dprob_ref"] = torch.sum(torch.cat(chunk_sums, dim=1), dim=1, keepdim=True) - prob = prob_tensor.expand(-1, n, -1) - ab = ref * prob * swish - dswiglu = ref * prob * c_input * sig * (1 + c_gate * (1 - sig)) - ref_d = torch.empty_like(c_full) ref_d.index_copy_(dim=1, index=dest_idx_ab, source=ab.index_select(dim=1, index=src_idx_n)) ref_d.index_copy_(dim=1, index=dest_idx_glu, source=dswiglu.index_select(dim=1, index=src_idx_n)) @@ -426,6 +441,9 @@ def check_ref_discrete_dswiglu( d_dtype=cfg["d_dtype"], sf_vec_size=cfg["sf_vec_size"], sf_dtype=cfg["sf_dtype"], + act_func=cfg.get("act_func", "dswiglu"), + situ_beta1=cfg.get("situ_beta1", 4.0), + situ_beta2=cfg.get("situ_beta2", 25.0), ) torch.cuda.synchronize() diff --git a/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_utils.py b/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_utils.py index dc5fac3ce..f7e9da794 100644 --- a/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_utils.py +++ b/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_utils.py @@ -367,6 +367,9 @@ def run_discrete_grouped_gemm_ref( generate_amax: bool = False, c_dtype: torch.dtype = torch.bfloat16, d_dtype: torch.dtype = torch.float32, + act_func: str = "swiglu", + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, ) -> Dict[str, torch.Tensor]: """Run reference implementation for discrete grouped GEMM GLU (SwiGLU). @@ -432,8 +435,12 @@ def run_discrete_grouped_gemm_ref( ref_gate = ref.index_select(1, gate_idx) ref_up = ref.index_select(1, up_idx) - # SwiGLU: up * (gate * sigmoid(gate)) - ref_gate = ref_gate * torch.sigmoid(ref_gate) + if act_func == "situglu": + ref_gate = situ_beta1 * torch.tanh(ref_gate / situ_beta1) * torch.sigmoid(ref_gate) + ref_up = situ_beta2 * torch.tanh(ref_up / situ_beta2) + else: + # SwiGLU: up * (gate * sigmoid(gate)) + ref_gate = ref_gate * torch.sigmoid(ref_gate) ref_after_swiglu = ref_up * ref_gate # Step 4: Apply prob @@ -496,6 +503,9 @@ def check_ref_discrete_grouped_gemm( generate_amax=(outputs.get("amax_tensor") is not None), c_dtype=cfg["c_dtype"], d_dtype=cfg["d_dtype"], + act_func=cfg.get("act_func", "swiglu"), + situ_beta1=cfg.get("situ_beta1", 4.0), + situ_beta2=cfg.get("situ_beta2", 25.0), ) torch.cuda.synchronize() diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py index 770e64af7..50851945f 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py @@ -755,6 +755,10 @@ def _test_grouped_gemm_dglu_dense_wrapper( omit_prob=False, use_single_group_runtime_offsets=False, sf_fp8_dtype_override=None, + act_func="dswiglu", + situ_beta1=4.0, + situ_beta2=25.0, + skip_reference=False, ): try: from cudnn import grouped_gemm_dglu_wrapper_sm100 @@ -778,6 +782,9 @@ def _test_grouped_gemm_dglu_dense_wrapper( b_major=b_major, ) cfg = _apply_grouped_gemm_cfg_overrides(cfg, cfg_overrides) + cfg["act_func"] = act_func + cfg["situ_beta1"] = situ_beta1 + cfg["situ_beta2"] = situ_beta2 stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) @@ -845,7 +852,9 @@ def _test_grouped_gemm_dglu_dense_wrapper( vector_f32=cfg["vector_f32"], m_aligned=cfg["m_aligned"], discrete_col_sfd=cfg["discrete_col_sfd"], - act_func="dswiglu", + act_func=act_func, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, use_dynamic_sched=use_dynamic_sched, use_single_group_runtime_offsets=use_single_group_runtime_offsets, current_stream=stream, @@ -854,10 +863,35 @@ def _test_grouped_gemm_dglu_dense_wrapper( pytest.skip(f"Unsupported testcase: {e}") torch.cuda.synchronize() - check_ref_grouped_gemm_dswiglu(inputs, wrapper_outputs, cfg, skip_ref=cfg["skip_ref"]) + if not skip_reference: + check_ref_grouped_gemm_dswiglu(inputs, wrapper_outputs, cfg, skip_ref=cfg["skip_ref"]) return inputs, wrapper_outputs, cfg +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_dglu_dense_wrapper_dsituglu_mxfp8(request): + """Smoke-test the dense MXFP8 compile and execute path for dSiTU-GLU.""" + + _test_grouped_gemm_dglu_dense_wrapper( + ab_dtype=torch.float8_e4m3fn, + c_dtype=torch.bfloat16, + d_dtype=torch.float8_e4m3fn, + b_major="k", + cd_major="n", + acc_dtype=torch.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + sf_dtype=torch.float8_e8m0fnu, + vector_f32=False, + discrete_col_sfd=False, + request=request, + act_func="dsituglu", + skip_reference=True, + ) + + @pytest.mark.L0 @torch_fork_set_rng(seed=2) @pytest.mark.parametrize( @@ -1343,6 +1377,8 @@ def _test_grouped_gemm_dglu_discrete_wrapper( b_major="k", generate_dbias=False, use_dynamic_sched=False, + situ_beta1=4.0, + situ_beta2=25.0, ): try: from cudnn import grouped_gemm_dglu_wrapper_sm100 @@ -1366,6 +1402,8 @@ def _test_grouped_gemm_dglu_discrete_wrapper( act_func, b_major=b_major, ) + cfg["situ_beta1"] = situ_beta1 + cfg["situ_beta2"] = situ_beta2 stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) @@ -1411,6 +1449,8 @@ def _test_grouped_gemm_dglu_discrete_wrapper( m_aligned=cfg["m_aligned"], discrete_col_sfd=cfg["discrete_col_sfd"], act_func=cfg["act_func"], + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, use_dynamic_sched=use_dynamic_sched, current_stream=stream, ) @@ -1421,6 +1461,40 @@ def _test_grouped_gemm_dglu_discrete_wrapper( check_ref_discrete_dswiglu(inputs, outputs, cfg, skip_ref=cfg["skip_ref"]) +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize( + ("ab_dtype", "d_dtype"), + [ + (torch.float4_e2m1fn_x2, torch.bfloat16), + (torch.float8_e4m3fn, torch.float8_e4m3fn), + ], + ids=["mxfp4", "mxfp8"], +) +@pytest.mark.parametrize(("situ_beta1", "situ_beta2"), [(4.0, 25.0), (2.0, 8.0)]) +def test_grouped_gemm_dglu_discrete_wrapper_dsituglu(ab_dtype, d_dtype, situ_beta1, situ_beta2, request): + """Exercise dSiTU-GLU with the inherited MXFP4 and MXFP8 layouts.""" + + _test_grouped_gemm_dglu_discrete_wrapper( + ab_dtype=ab_dtype, + c_dtype=torch.bfloat16, + d_dtype=d_dtype, + cd_major="n", + acc_dtype=torch.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + sf_dtype=torch.float8_e8m0fnu, + vector_f32=False, + discrete_col_sfd=False, + act_func="dsituglu", + request=request, + b_major="k", + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, + ) + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) def test_grouped_gemm_dglu_discrete_compile_execute_with_dbias(request): diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py index 9d0a9e73a..95a48fb46 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py @@ -682,6 +682,10 @@ def _test_grouped_gemm_glu_dense_wrapper( omit_prob=False, use_single_group_runtime_offsets=False, sf_fp8_dtype_override=None, + act_func="swiglu", + situ_beta1=4.0, + situ_beta2=25.0, + skip_reference=False, ): try: from cudnn import grouped_gemm_glu_wrapper_sm100 @@ -705,6 +709,9 @@ def _test_grouped_gemm_glu_dense_wrapper( enable_bias=enable_bias, ) cfg = _apply_grouped_gemm_cfg_overrides(cfg, cfg_overrides) + cfg["act_func"] = act_func + cfg["situ_beta1"] = situ_beta1 + cfg["situ_beta2"] = situ_beta2 stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) @@ -755,7 +762,9 @@ def _test_grouped_gemm_glu_dense_wrapper( vector_f32=cfg["vector_f32"], m_aligned=cfg["m_aligned"], discrete_col_sfd=cfg["discrete_col_sfd"], - act_func="swiglu", + act_func=act_func, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, use_dynamic_sched=use_dynamic_sched, use_single_group_runtime_offsets=use_single_group_runtime_offsets, current_stream=stream, @@ -763,10 +772,34 @@ def _test_grouped_gemm_glu_dense_wrapper( except (ValueError, NotImplementedError) as e: pytest.skip(f"Unsupported testcase: {e}") - check_ref_grouped_gemm_swiglu(inputs, outputs, cfg, skip_ref=cfg["skip_ref"]) + if not skip_reference: + check_ref_grouped_gemm_swiglu(inputs, outputs, cfg, skip_ref=cfg["skip_ref"]) return inputs, outputs, cfg +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_dense_wrapper_situglu_mxfp8(request): + """Smoke-test the dense MXFP8 compile and execute path for SiTU-GLU.""" + + _test_grouped_gemm_glu_dense_wrapper( + ab_dtype=torch.float8_e4m3fn, + c_dtype=torch.bfloat16, + d_dtype=torch.float8_e4m3fn, + cd_major="n", + acc_dtype=torch.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + sf_dtype=torch.float8_e8m0fnu, + vector_f32=False, + discrete_col_sfd=False, + request=request, + act_func="situglu", + skip_reference=True, + ) + + @pytest.mark.L0 @torch_fork_set_rng(seed=1) @pytest.mark.parametrize( @@ -1233,6 +1266,8 @@ def _test_grouped_gemm_glu_discrete_wrapper( b_major="k", enable_bias=False, use_dynamic_sched=False, + situ_beta1=4.0, + situ_beta2=25.0, ): try: from cudnn import grouped_gemm_glu_wrapper_sm100 @@ -1256,6 +1291,8 @@ def _test_grouped_gemm_glu_discrete_wrapper( act_func, b_major=b_major, ) + cfg["situ_beta1"] = situ_beta1 + cfg["situ_beta2"] = situ_beta2 stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) @@ -1298,6 +1335,8 @@ def _test_grouped_gemm_glu_discrete_wrapper( m_aligned=cfg["m_aligned"], discrete_col_sfd=cfg["discrete_col_sfd"], act_func=cfg["act_func"], + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, use_dynamic_sched=use_dynamic_sched, current_stream=stream, ) @@ -1307,6 +1346,40 @@ def _test_grouped_gemm_glu_discrete_wrapper( check_ref_discrete_grouped_gemm(inputs, outputs, cfg, skip_ref=cfg["skip_ref"]) +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize( + ("ab_dtype", "d_dtype"), + [ + (torch.float4_e2m1fn_x2, torch.bfloat16), + (torch.float8_e4m3fn, torch.float8_e4m3fn), + ], + ids=["mxfp4", "mxfp8"], +) +@pytest.mark.parametrize(("situ_beta1", "situ_beta2"), [(4.0, 25.0), (2.0, 8.0)]) +def test_grouped_gemm_glu_discrete_wrapper_situglu(ab_dtype, d_dtype, situ_beta1, situ_beta2, request): + """Exercise SiTU-GLU with the inherited MXFP4 and MXFP8 layouts.""" + + _test_grouped_gemm_glu_discrete_wrapper( + ab_dtype=ab_dtype, + c_dtype=torch.bfloat16, + d_dtype=d_dtype, + cd_major="n", + acc_dtype=torch.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + sf_dtype=torch.float8_e8m0fnu, + vector_f32=False, + discrete_col_sfd=False, + act_func="situglu", + request=request, + b_major="k", + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, + ) + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) def test_grouped_gemm_glu_discrete_compile_execute_with_bias(request): From 2fe36f4ad4450efb30d9ec355a7838b258e94c40 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 18 Aug 2026 18:12:10 +0800 Subject: [PATCH 2/9] [PyTorch] Optimize K3 SiTU-GLU activation Signed-off-by: Harry Zhou --- .../cutedsl/grouped/dglu/_blockscaled_api.py | 1 + ...moe_blockscaled_grouped_gemm_dglu_dbias.py | 38 +++++++++++++++++-- .../cutedsl/grouped/glu/_blockscaled_api.py | 15 ++++++++ python/cudnn/gemm/cutedsl/grouped/glu/api.py | 6 +++ .../moe_blockscaled_grouped_gemm_glu_bias.py | 12 +++++- 5 files changed, 68 insertions(+), 4 deletions(-) diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py index 55790da3f..a10ef7944 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py @@ -762,6 +762,7 @@ def compile(self) -> None: weight_mode=self.weight_mode, act_func=self.act_func, use_dynamic_sched=self.use_dynamic_sched, + **({"situ_beta1": self.situ_beta1} if not self._is_rubin_kernel else {}), **rubin_single_group_offsets_kwarg(self._is_rubin_kernel, self.use_single_group_runtime_offsets), # Only the Rubin kernel accepts sf_fp8_dtype_override, and check_support # rejects "e5m3" unless _is_rubin_kernel -- the same flag that selected diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py index 56ef5757a..4135af08e 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py @@ -259,6 +259,7 @@ def __init__( weight_mode: MoEWeightMode = MoEWeightMode.DISCRETE, use_dynamic_sched: bool = False, act_func: str = "dswiglu", + situ_beta1: float = 4.0, use_single_group_runtime_offsets: bool = False, ): """Initializes the configuration for a Blackwell blockscaled grouped GEMM dGLU kernel. @@ -374,6 +375,7 @@ def __init__( self.weight_mode = weight_mode self.act_func = act_func + self.situ_beta1 = float(situ_beta1) def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -1904,11 +1906,41 @@ def dsituglu( up = up_vec[i].to(cutlass.Float32) * beta_val gate_tanh = cute.math.tanh(gate / beta1, fastmath=True) up_tanh = cute.math.tanh(up / beta2, fastmath=True) - sigmoid = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True)) + if cutlass.const_expr(self.situ_beta1 == 4.0): + # For a = tanh(gate / 4), sigmoid(gate) = 1/2 + a / (1 + a^2). + gate_tanh_sq = gate_tanh * gate_tanh + gate_tanh_denom_rcp = cute.arch.rcp_approx( + cutlass.Float32(1.0) + gate_tanh_sq + ) + sigmoid = ( + cutlass.Float32(0.5) + gate_tanh * gate_tanh_denom_rcp + ) + else: + sigmoid = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True) + ) gate_value = beta1 * gate_tanh * sigmoid up_value = beta2 * up_tanh - gate_grad = (cutlass.Float32(1.0) - gate_tanh * gate_tanh) * sigmoid - gate_grad = gate_grad + beta1 * gate_tanh * sigmoid * (cutlass.Float32(1.0) - sigmoid) + if cutlass.const_expr(self.situ_beta1 == 4.0): + # d[4*a*sigmoid(gate)]/dgate, expressed with the same reciprocal. + gate_grad = (cutlass.Float32(1.0) - gate_tanh_sq) * ( + cutlass.Float32(0.5) + + cutlass.Float32(2.0) + * gate_tanh + * gate_tanh_denom_rcp + * gate_tanh_denom_rcp + ) + else: + gate_grad = ( + cutlass.Float32(1.0) - gate_tanh * gate_tanh + ) * sigmoid + gate_grad = ( + gate_grad + + beta1 + * gate_tanh + * sigmoid + * (cutlass.Float32(1.0) - sigmoid) + ) up_grad = cutlass.Float32(1.0) - up_tanh * up_tanh activation_grad = grad if cutlass.const_expr(self.has_prob): diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py index d55454ca9..af5827543 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py @@ -128,6 +128,7 @@ def __init__( m_aligned: int = 256, discrete_col_sfd: bool = False, act_func: str = "swiglu", + situ_beta1: float = 4.0, b_major: str = "k", use_dynamic_sched: bool = False, use_single_group_runtime_offsets: bool = False, @@ -167,6 +168,8 @@ def __init__( :param m_aligned: Alignment for group M dimension :param discrete_col_sfd: Generate discrete col-major scale factor tensor :param act_func: Activation function, one of "swiglu", "geglu", or "situglu" + :param situ_beta1: Compile-time gate tanh scale for SiTU-GLU. The default + ``4.0`` enables an exact sigmoid-reuse specialization. :param b_major: Major dimension for B tensor, one of "k" or "n" :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing """ @@ -248,6 +251,12 @@ def __init__( self.m_aligned = m_aligned self.discrete_col_sfd = discrete_col_sfd self.act_func = act_func + self.situ_beta1 = float(situ_beta1) + if self.act_func == "situglu": + self._value_error_if( + not math.isfinite(self.situ_beta1) or self.situ_beta1 <= 0.0, + f"situ_beta1 must be finite and positive, got {self.situ_beta1}", + ) if self.weight_mode == MoEWeightMode.DENSE: self.b_major = b_major # stored for both modes @@ -678,6 +687,7 @@ def compile(self) -> None: act_func=self.act_func, enable_bias=self._has_bias, use_dynamic_sched=self.use_dynamic_sched, + **({"situ_beta1": self.situ_beta1} if not self._is_rubin_kernel else {}), **rubin_single_group_offsets_kwarg(self._is_rubin_kernel, self.use_single_group_runtime_offsets), # Only the Rubin kernel accepts sf_fp8_dtype_override, and check_support # rejects "e5m3" unless _is_rubin_kernel -- the same flag that selected @@ -1309,6 +1319,11 @@ def execute( not math.isfinite(situ_beta2) or situ_beta2 <= 0.0, f"situ_beta2 must be finite and positive, got {situ_beta2}", ) + self._value_error_if( + float(situ_beta1) != self.situ_beta1, + "situ_beta1 is specialized at compile time; construct and compile " + f"the API with situ_beta1={situ_beta1}", + ) _reject_unsupported_rubin_glu_tune_params( self._is_rubin_kernel, geglu_alpha, diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/api.py b/python/cudnn/gemm/cutedsl/grouped/glu/api.py index 9df4063d0..2a98523c4 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/api.py @@ -190,6 +190,7 @@ def __init__( m_aligned: int = 256, discrete_col_sfd: bool = False, act_func: str = "swiglu", + situ_beta1: float = 4.0, b_major: str = "k", use_dynamic_sched: bool = False, use_single_group_runtime_offsets: bool = False, @@ -659,6 +660,7 @@ def dynamic_m_tensor_signature( return static_shape_suffix, stride_signature, tensor.dtype use_full_dynamic = is_dense and os.environ.get("CUDNN_FE_GROUPED_GEMM_DYNAMIC_MNKL", "1") != "0" + situ_beta1_cache_signature = float(situ_beta1) if act_func == "situglu" else None device_type = get_device_type() @@ -667,6 +669,7 @@ def dynamic_m_tensor_signature( device_type, weight_mode, act_func, + situ_beta1_cache_signature, use_full_dynamic, a_tensor.shape[1:] if not use_full_dynamic else None, b_tensor.shape[2] if use_full_dynamic else tuple(b_tensor.shape), @@ -711,6 +714,7 @@ def dynamic_m_tensor_signature( device_type, weight_mode, act_func, + situ_beta1_cache_signature, a_tensor.shape[1:], stride_order(a_tensor), a_tensor.dtype, @@ -784,6 +788,7 @@ def dynamic_m_tensor_signature( m_aligned=m_aligned, discrete_col_sfd=discrete_col_sfd, act_func=act_func, + situ_beta1=situ_beta1, use_dynamic_sched=use_dynamic_sched, use_single_group_runtime_offsets=use_single_group_runtime_offsets, ) @@ -814,6 +819,7 @@ def dynamic_m_tensor_signature( m_aligned=m_aligned, discrete_col_sfd=discrete_col_sfd, act_func=act_func, + situ_beta1=situ_beta1, b_major=b_major, use_dynamic_sched=use_dynamic_sched, use_single_group_runtime_offsets=use_single_group_runtime_offsets, diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py b/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py index ae03e6c4f..9b65756a1 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py @@ -263,6 +263,7 @@ def __init__( weight_mode: MoEWeightMode = MoEWeightMode.DISCRETE, use_dynamic_sched: bool = False, act_func: str = "swiglu", + situ_beta1: float = 4.0, enable_bias: bool = False, use_single_group_runtime_offsets: bool = False, ): @@ -381,6 +382,7 @@ def __init__( self.num_epilog_warps = len(self.epilog_warp_id) self.act_func = act_func + self.situ_beta1 = float(situ_beta1) if act_func not in ["swiglu", "geglu", "situglu"]: raise ValueError(f"Invalid activation function: {act_func}") @@ -1578,7 +1580,15 @@ def situglu_act( up = acc_vec_up[i] gate_tanh = cute.math.tanh(gate / beta1, fastmath=True) up_tanh = cute.math.tanh(up / beta2, fastmath=True) - sigmoid = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True)) + if cutlass.const_expr(self.situ_beta1 == 4.0): + # For a = tanh(gate / 4), sigmoid(gate) = 1/2 + a / (1 + a^2). + sigmoid = cutlass.Float32(0.5) + gate_tanh * cute.arch.rcp_approx( + cutlass.Float32(1.0) + gate_tanh * gate_tanh + ) + else: + sigmoid = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True) + ) tCompute[i] = beta1 * gate_tanh * sigmoid * beta2 * up_tanh if cutlass.const_expr(self.has_prob): tCompute[i] = tCompute[i] * mProb From bf4cc07005f5f7e2bf6f0e9eeeeefa35a5db9499 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 18 Aug 2026 19:17:22 +0800 Subject: [PATCH 3/9] [CuTeDSL] Fix SiTU-GLU documentation Signed-off-by: Harry Zhou --- docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md | 4 ++-- docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md | 9 +++++---- python/cudnn/gemm/cutedsl/grouped/glu/api.py | 3 ++- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md index d3db216a9..56f60f8bf 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md @@ -146,8 +146,8 @@ $$ \frac{\partial T_u}{\partial U}=1-\tanh^2(U/\beta_2), $$ -and returns `R * prob * T_u * dT_g/dG` and -`R * prob * T_g * dT_u/dU`. `dprob` accumulates `R * T_g * T_u`. +and returns `ref * prob * T_u * dT_g/dG` and +`ref * prob * T_g * dT_u/dU`. `dprob` accumulates `ref * T_g * T_u`. The beta values are compile-time specialization values and therefore belong to the dGLU compiled-kernel cache key. diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md index 1456193bf..83cafd4bb 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md @@ -100,8 +100,7 @@ The block-scaled backend performs: ### Equations -For SiTU-GLU, with gate branch `G`, up branch `U`, `situ_beta1=4`, and -`situ_beta2=25`, the fused epilogue computes +For SiTU-GLU, with gate branch `G` and up branch `U`, the fused epilogue computes $$ D = \mathrm{prob}\, @@ -109,8 +108,10 @@ D = \mathrm{prob}\, \left[\beta_2\tanh(U/\beta_2)\right]. $$ -The beta values are runtime FP32 scalars and do not create new compiled-kernel -cache entries. +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)` diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/api.py b/python/cudnn/gemm/cutedsl/grouped/glu/api.py index 2a98523c4..96f366b29 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/api.py @@ -502,7 +502,8 @@ def _grouped_gemm_glu_block_scaled_call(call: GluCall) -> TupleDict: parameter, intentionally not part of the cache key. Ignored when ``act_func == "swiglu"``. situ_beta1: Positive finite gate tanh scale for SiTU-GLU. Default - ``4.0``. Runtime parameter, intentionally not part of the cache key. + ``4.0``. This value specializes the compiled kernel and is part of + the cache key. situ_beta2: Positive finite up-branch tanh scale for SiTU-GLU. Default ``25.0``. Runtime parameter, intentionally not part of the cache key. use_dynamic_sched: Enable dynamic tile scheduling for load balancing From 230ff62249ef584ead2456b5c97d2590c4d4417c Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 18 Aug 2026 20:55:02 +0800 Subject: [PATCH 4/9] [CuTeDSL] Optimize SiTU-GLU activation math Signed-off-by: Harry Zhou --- ...moe_blockscaled_grouped_gemm_dglu_dbias.py | 135 +++++++++++++----- .../moe_blockscaled_grouped_gemm_glu_bias.py | 17 ++- 2 files changed, 111 insertions(+), 41 deletions(-) diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py index 4135af08e..4e28bc34e 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py @@ -1893,54 +1893,125 @@ def dsituglu( mProb: cute.Tensor, beta_val: Float32, square_alpha: Float32, - beta1: Float32, - beta2: Float32, + beta1: cutlass.Constexpr, + beta2: cutlass.Constexpr, dprob_swiglu: Optional[cute.Tensor] = None, ): dgate_vec = cute.make_rmem_tensor(acc_vec.shape, cutlass.Float32) dup_vec = cute.make_rmem_tensor(acc_vec.shape, cutlass.Float32) + beta1_f32 = cutlass.Float32(beta1) + beta2_f32 = cutlass.Float32(beta2) + beta1_rcp = cutlass.Float32(1.0 / beta1) + beta2_rcp = cutlass.Float32(1.0 / beta2) + + if cutlass.const_expr(self.situ_beta1 == 4.0): + fmul2 = partial(cute.arch.mul_packed_f32x2, rnd="rn", ftz=False) + fadd2 = partial(cute.arch.add_packed_f32x2, rnd="rn", ftz=False) + square_alpha2 = (square_alpha, square_alpha) + beta_val2 = (beta_val, beta_val) + beta1_2 = (beta1_f32, beta1_f32) + beta2_2 = (beta2_f32, beta2_f32) + beta1_rcp2 = (beta1_rcp, beta1_rcp) + beta2_rcp2 = (beta2_rcp, beta2_rcp) + mprob2 = (mProb, mProb) + ones2 = (cutlass.Float32(1.0), cutlass.Float32(1.0)) + halves2 = (cutlass.Float32(0.5), cutlass.Float32(0.5)) + twos2 = (cutlass.Float32(2.0), cutlass.Float32(2.0)) + + for i in cutlass.range(0, cute.size(acc_vec), 2, unroll_full=True): + grad = fmul2((acc_vec[i], acc_vec[i + 1]), square_alpha2) + gate = fmul2( + ( + gate_vec[i].to(cutlass.Float32), + gate_vec[i + 1].to(cutlass.Float32), + ), + beta_val2, + ) + up = fmul2( + ( + up_vec[i].to(cutlass.Float32), + up_vec[i + 1].to(cutlass.Float32), + ), + beta_val2, + ) + gate_scaled = fmul2(gate, beta1_rcp2) + up_scaled = fmul2(up, beta2_rcp2) + gate_tanh = ( + cute.math.tanh(gate_scaled[0], fastmath=True), + cute.math.tanh(gate_scaled[1], fastmath=True), + ) + up_tanh = ( + cute.math.tanh(up_scaled[0], fastmath=True), + cute.math.tanh(up_scaled[1], fastmath=True), + ) + gate_tanh_sq = fmul2(gate_tanh, gate_tanh) + gate_tanh_denom = fadd2(ones2, gate_tanh_sq) + gate_tanh_denom_rcp = ( + cute.arch.rcp_approx(gate_tanh_denom[0]), + cute.arch.rcp_approx(gate_tanh_denom[1]), + ) + sigmoid = fadd2( + halves2, + fmul2(gate_tanh, gate_tanh_denom_rcp), + ) + gate_value = fmul2(fmul2(beta1_2, gate_tanh), sigmoid) + up_value = fmul2(beta2_2, up_tanh) + gate_tanh_denom_rcp_sq = fmul2( + gate_tanh_denom_rcp, + gate_tanh_denom_rcp, + ) + gate_grad_inner = fadd2( + halves2, + fmul2( + twos2, + fmul2(gate_tanh, gate_tanh_denom_rcp_sq), + ), + ) + one_minus_gate_tanh_sq = fadd2( + ones2, + (-gate_tanh_sq[0], -gate_tanh_sq[1]), + ) + gate_grad = fmul2(one_minus_gate_tanh_sq, gate_grad_inner) + up_tanh_sq = fmul2(up_tanh, up_tanh) + up_grad = fadd2(ones2, (-up_tanh_sq[0], -up_tanh_sq[1])) + activation_grad = grad + if cutlass.const_expr(self.has_prob): + activation_grad = fmul2(grad, mprob2) + dgate = fmul2(fmul2(activation_grad, up_value), gate_grad) + dup = fmul2(fmul2(activation_grad, gate_value), up_grad) + dgate_vec[i], dgate_vec[i + 1] = dgate + dup_vec[i], dup_vec[i + 1] = dup + if cutlass.const_expr(self.generate_dprob): + dprob = fmul2(fmul2(grad, gate_value), up_value) + dprob_swiglu[i], dprob_swiglu[i + 1] = dprob + + if cutlass.const_expr(self.generate_dprob): + dprob_swiglu = dprob_swiglu.load() + return dgate_vec.load(), dup_vec.load(), dprob_swiglu for i in cutlass.range_constexpr(cute.size(acc_vec)): grad = acc_vec[i] * square_alpha gate = gate_vec[i].to(cutlass.Float32) * beta_val up = up_vec[i].to(cutlass.Float32) * beta_val - gate_tanh = cute.math.tanh(gate / beta1, fastmath=True) - up_tanh = cute.math.tanh(up / beta2, fastmath=True) + gate_tanh = cute.math.tanh(gate * beta1_rcp, fastmath=True) + up_tanh = cute.math.tanh(up * beta2_rcp, fastmath=True) if cutlass.const_expr(self.situ_beta1 == 4.0): # For a = tanh(gate / 4), sigmoid(gate) = 1/2 + a / (1 + a^2). gate_tanh_sq = gate_tanh * gate_tanh - gate_tanh_denom_rcp = cute.arch.rcp_approx( - cutlass.Float32(1.0) + gate_tanh_sq - ) - sigmoid = ( - cutlass.Float32(0.5) + gate_tanh * gate_tanh_denom_rcp - ) + gate_tanh_denom_rcp = cute.arch.rcp_approx(cutlass.Float32(1.0) + gate_tanh_sq) + sigmoid = cutlass.Float32(0.5) + gate_tanh * gate_tanh_denom_rcp else: - sigmoid = cute.arch.rcp_approx( - cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True) - ) - gate_value = beta1 * gate_tanh * sigmoid - up_value = beta2 * up_tanh + sigmoid = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True)) + gate_value = beta1_f32 * gate_tanh * sigmoid + up_value = beta2_f32 * up_tanh if cutlass.const_expr(self.situ_beta1 == 4.0): # d[4*a*sigmoid(gate)]/dgate, expressed with the same reciprocal. gate_grad = (cutlass.Float32(1.0) - gate_tanh_sq) * ( - cutlass.Float32(0.5) - + cutlass.Float32(2.0) - * gate_tanh - * gate_tanh_denom_rcp - * gate_tanh_denom_rcp + cutlass.Float32(0.5) + cutlass.Float32(2.0) * gate_tanh * gate_tanh_denom_rcp * gate_tanh_denom_rcp ) else: - gate_grad = ( - cutlass.Float32(1.0) - gate_tanh * gate_tanh - ) * sigmoid - gate_grad = ( - gate_grad - + beta1 - * gate_tanh - * sigmoid - * (cutlass.Float32(1.0) - sigmoid) - ) + gate_grad = (cutlass.Float32(1.0) - gate_tanh * gate_tanh) * sigmoid + gate_grad = gate_grad + beta1_f32 * gate_tanh * sigmoid * (cutlass.Float32(1.0) - sigmoid) up_grad = cutlass.Float32(1.0) - up_tanh * up_tanh activation_grad = grad if cutlass.const_expr(self.has_prob): @@ -3128,8 +3199,8 @@ def kernel( mProb, beta_val, square_alpha, - cutlass.Float32(situ_beta1), - cutlass.Float32(situ_beta2), + situ_beta1, + situ_beta2, dprob_swiglu, ) diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py b/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py index 9b65756a1..6813e05fa 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_bias.py @@ -1575,21 +1575,20 @@ def situglu_act( beta1: cutlass.Float32, beta2: cutlass.Float32, ): + beta1_rcp = cutlass.Float32(1.0 / self.situ_beta1) + beta2_rcp = cute.arch.rcp_approx(beta2) + beta_product = beta1 * beta2 for i in cutlass.range_constexpr(cute.size(tCompute)): gate = acc_vec_gate[i] up = acc_vec_up[i] - gate_tanh = cute.math.tanh(gate / beta1, fastmath=True) - up_tanh = cute.math.tanh(up / beta2, fastmath=True) + gate_tanh = cute.math.tanh(gate * beta1_rcp, fastmath=True) + up_tanh = cute.math.tanh(up * beta2_rcp, fastmath=True) if cutlass.const_expr(self.situ_beta1 == 4.0): # For a = tanh(gate / 4), sigmoid(gate) = 1/2 + a / (1 + a^2). - sigmoid = cutlass.Float32(0.5) + gate_tanh * cute.arch.rcp_approx( - cutlass.Float32(1.0) + gate_tanh * gate_tanh - ) + sigmoid = cutlass.Float32(0.5) + gate_tanh * cute.arch.rcp_approx(cutlass.Float32(1.0) + gate_tanh * gate_tanh) else: - sigmoid = cute.arch.rcp_approx( - cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True) - ) - tCompute[i] = beta1 * gate_tanh * sigmoid * beta2 * up_tanh + sigmoid = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True)) + tCompute[i] = beta_product * gate_tanh * sigmoid * up_tanh if cutlass.const_expr(self.has_prob): tCompute[i] = tCompute[i] * mProb From 79cabc31dac6e85329ac658e6237f7547fa3ce4f Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 18 Aug 2026 21:13:13 +0800 Subject: [PATCH 5/9] [CuTeDSL] Test NVFP4 SiTU-GLU paths Signed-off-by: Harry Zhou --- .../grouped_gemm/test_grouped_gemm_dglu.py | 17 +++++++++-------- .../grouped_gemm/test_grouped_gemm_glu.py | 17 +++++++++-------- 2 files changed, 18 insertions(+), 16 deletions(-) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py index 50851945f..0872ce5c5 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py @@ -1464,16 +1464,17 @@ def _test_grouped_gemm_dglu_discrete_wrapper( @pytest.mark.L0 @torch_fork_set_rng(seed=0) @pytest.mark.parametrize( - ("ab_dtype", "d_dtype"), + ("ab_dtype", "d_dtype", "sf_dtype", "sf_vec_size"), [ - (torch.float4_e2m1fn_x2, torch.bfloat16), - (torch.float8_e4m3fn, torch.float8_e4m3fn), + (torch.float4_e2m1fn_x2, torch.bfloat16, torch.float8_e8m0fnu, 32), + (torch.float8_e4m3fn, torch.float8_e4m3fn, torch.float8_e8m0fnu, 32), + (torch.float4_e2m1fn_x2, torch.bfloat16, torch.float8_e4m3fn, 16), ], - ids=["mxfp4", "mxfp8"], + ids=["mxfp4", "mxfp8", "nvfp4"], ) @pytest.mark.parametrize(("situ_beta1", "situ_beta2"), [(4.0, 25.0), (2.0, 8.0)]) -def test_grouped_gemm_dglu_discrete_wrapper_dsituglu(ab_dtype, d_dtype, situ_beta1, situ_beta2, request): - """Exercise dSiTU-GLU with the inherited MXFP4 and MXFP8 layouts.""" +def test_grouped_gemm_dglu_discrete_wrapper_dsituglu(ab_dtype, d_dtype, sf_dtype, sf_vec_size, situ_beta1, situ_beta2, request): + """Exercise dSiTU-GLU with the inherited MXFP4, MXFP8, and NVFP4 layouts.""" _test_grouped_gemm_dglu_discrete_wrapper( ab_dtype=ab_dtype, @@ -1483,8 +1484,8 @@ def test_grouped_gemm_dglu_discrete_wrapper_dsituglu(ab_dtype, d_dtype, situ_bet acc_dtype=torch.float32, mma_tiler_mn=(256, 256), cluster_shape_mn=(2, 1), - sf_vec_size=32, - sf_dtype=torch.float8_e8m0fnu, + sf_vec_size=sf_vec_size, + sf_dtype=sf_dtype, vector_f32=False, discrete_col_sfd=False, act_func="dsituglu", diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py index 95a48fb46..96e875586 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py @@ -1349,16 +1349,17 @@ def _test_grouped_gemm_glu_discrete_wrapper( @pytest.mark.L0 @torch_fork_set_rng(seed=0) @pytest.mark.parametrize( - ("ab_dtype", "d_dtype"), + ("ab_dtype", "d_dtype", "sf_dtype", "sf_vec_size"), [ - (torch.float4_e2m1fn_x2, torch.bfloat16), - (torch.float8_e4m3fn, torch.float8_e4m3fn), + (torch.float4_e2m1fn_x2, torch.bfloat16, torch.float8_e8m0fnu, 32), + (torch.float8_e4m3fn, torch.float8_e4m3fn, torch.float8_e8m0fnu, 32), + (torch.float4_e2m1fn_x2, torch.bfloat16, torch.float8_e4m3fn, 16), ], - ids=["mxfp4", "mxfp8"], + ids=["mxfp4", "mxfp8", "nvfp4"], ) @pytest.mark.parametrize(("situ_beta1", "situ_beta2"), [(4.0, 25.0), (2.0, 8.0)]) -def test_grouped_gemm_glu_discrete_wrapper_situglu(ab_dtype, d_dtype, situ_beta1, situ_beta2, request): - """Exercise SiTU-GLU with the inherited MXFP4 and MXFP8 layouts.""" +def test_grouped_gemm_glu_discrete_wrapper_situglu(ab_dtype, d_dtype, sf_dtype, sf_vec_size, situ_beta1, situ_beta2, request): + """Exercise SiTU-GLU with the inherited MXFP4, MXFP8, and NVFP4 layouts.""" _test_grouped_gemm_glu_discrete_wrapper( ab_dtype=ab_dtype, @@ -1368,8 +1369,8 @@ def test_grouped_gemm_glu_discrete_wrapper_situglu(ab_dtype, d_dtype, situ_beta1 acc_dtype=torch.float32, mma_tiler_mn=(256, 256), cluster_shape_mn=(2, 1), - sf_vec_size=32, - sf_dtype=torch.float8_e8m0fnu, + sf_vec_size=sf_vec_size, + sf_dtype=sf_dtype, vector_f32=False, discrete_col_sfd=False, act_func="situglu", From 1f2ef252cd1848c5915ec0068a91f80263eeb8b2 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Wed, 19 Aug 2026 15:44:27 +0800 Subject: [PATCH 6/9] [CuTeDSL] Remove unreachable dSiTU-GLU branches Signed-off-by: Harry Zhou --- .../gemm_fusions/grouped_gemm_dglu.md | 4 +++- .../cutedsl/grouped/dglu/_blockscaled_api.py | 4 +++- python/cudnn/gemm/cutedsl/grouped/dglu/api.py | 4 +++- ...moe_blockscaled_grouped_gemm_dglu_dbias.py | 20 +++++-------------- .../grouped_gemm/test_grouped_gemm_dglu.py | 16 ++++++++++++--- 5 files changed, 27 insertions(+), 21 deletions(-) diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md index 56f60f8bf..e3d38bffa 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md @@ -518,7 +518,9 @@ 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), `"dgeglu"`, or block-scaled `"dsituglu"` diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py index a10ef7944..3d4634e81 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py @@ -177,7 +177,9 @@ def __init__( knob existed. ``"e5m3"`` requires Rubin and the NVFP4 recipe, and the scale tensors are still supplied as ``torch.float8_e4m3fn`` because torch has no e5m3 dtype -- only the CuTe element type is overridden. - :param vector_f32: Use vectorized f32 operations + :param vector_f32: Use vectorized f32 operations for dSwiGLU and dGeGLU. + K3-default dSiTU-GLU (``situ_beta1=4.0``) always uses its packed + FP32x2 specialization; non-default dSiTU-GLU uses scalar FP32. :param m_aligned: Alignment for group M dimension :param discrete_col_sfd: Generate discrete col-major scale factor tensor :param act_func: Activation function, one of "dswiglu", "dgeglu", or "dsituglu" diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/api.py b/python/cudnn/gemm/cutedsl/grouped/dglu/api.py index 405386c63..cec1473b8 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/api.py @@ -485,7 +485,9 @@ def _grouped_gemm_dglu_block_scaled_call(call: DgluCall) -> TupleDict: bit to widen the scale range; it is Rubin-only, requires the NVFP4 recipe, and the scale tensors are still passed as ``torch.float8_e4m3fn`` because torch has no e5m3 dtype. - vector_f32: Use vectorized f32 + vector_f32: Use vectorized f32 for dSwiGLU and dGeGLU. K3-default + dSiTU-GLU (``situ_beta1=4.0``) always uses its packed FP32x2 + specialization; non-default dSiTU-GLU uses scalar FP32. m_aligned: M alignment (must be 256) discrete_col_sfd: Generate discrete col-major scale factor tensor act_func: Activation function ("dswiglu", "dgeglu", or block-scaled "dsituglu") diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py index 4e28bc34e..621bd6134 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py @@ -1904,6 +1904,8 @@ def dsituglu( beta1_rcp = cutlass.Float32(1.0 / beta1) beta2_rcp = cutlass.Float32(1.0 / beta2) + # The K3-default specialization always uses packed FP32x2 arithmetic. This is an + # activation-specific optimization and is independent of the generic vectorized_f32 knob. if cutlass.const_expr(self.situ_beta1 == 4.0): fmul2 = partial(cute.arch.mul_packed_f32x2, rnd="rn", ftz=False) fadd2 = partial(cute.arch.add_packed_f32x2, rnd="rn", ftz=False) @@ -1995,23 +1997,11 @@ def dsituglu( up = up_vec[i].to(cutlass.Float32) * beta_val gate_tanh = cute.math.tanh(gate * beta1_rcp, fastmath=True) up_tanh = cute.math.tanh(up * beta2_rcp, fastmath=True) - if cutlass.const_expr(self.situ_beta1 == 4.0): - # For a = tanh(gate / 4), sigmoid(gate) = 1/2 + a / (1 + a^2). - gate_tanh_sq = gate_tanh * gate_tanh - gate_tanh_denom_rcp = cute.arch.rcp_approx(cutlass.Float32(1.0) + gate_tanh_sq) - sigmoid = cutlass.Float32(0.5) + gate_tanh * gate_tanh_denom_rcp - else: - sigmoid = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True)) + sigmoid = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True)) gate_value = beta1_f32 * gate_tanh * sigmoid up_value = beta2_f32 * up_tanh - if cutlass.const_expr(self.situ_beta1 == 4.0): - # d[4*a*sigmoid(gate)]/dgate, expressed with the same reciprocal. - gate_grad = (cutlass.Float32(1.0) - gate_tanh_sq) * ( - cutlass.Float32(0.5) + cutlass.Float32(2.0) * gate_tanh * gate_tanh_denom_rcp * gate_tanh_denom_rcp - ) - else: - gate_grad = (cutlass.Float32(1.0) - gate_tanh * gate_tanh) * sigmoid - gate_grad = gate_grad + beta1_f32 * gate_tanh * sigmoid * (cutlass.Float32(1.0) - sigmoid) + gate_grad = (cutlass.Float32(1.0) - gate_tanh * gate_tanh) * sigmoid + gate_grad = gate_grad + beta1_f32 * gate_tanh * sigmoid * (cutlass.Float32(1.0) - sigmoid) up_grad = cutlass.Float32(1.0) - up_tanh * up_tanh activation_grad = grad if cutlass.const_expr(self.has_prob): diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py index 0872ce5c5..cabccbb39 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py @@ -1473,8 +1473,18 @@ def _test_grouped_gemm_dglu_discrete_wrapper( ids=["mxfp4", "mxfp8", "nvfp4"], ) @pytest.mark.parametrize(("situ_beta1", "situ_beta2"), [(4.0, 25.0), (2.0, 8.0)]) -def test_grouped_gemm_dglu_discrete_wrapper_dsituglu(ab_dtype, d_dtype, sf_dtype, sf_vec_size, situ_beta1, situ_beta2, request): - """Exercise dSiTU-GLU with the inherited MXFP4, MXFP8, and NVFP4 layouts.""" +@pytest.mark.parametrize("vector_f32", [False, True], ids=["vector-f32-off", "vector-f32-on"]) +def test_grouped_gemm_dglu_discrete_wrapper_dsituglu( + ab_dtype, + d_dtype, + sf_dtype, + sf_vec_size, + situ_beta1, + situ_beta2, + vector_f32, + request, +): + """Exercise dSiTU-GLU layouts and its activation-specific vectorization policy.""" _test_grouped_gemm_dglu_discrete_wrapper( ab_dtype=ab_dtype, @@ -1486,7 +1496,7 @@ def test_grouped_gemm_dglu_discrete_wrapper_dsituglu(ab_dtype, d_dtype, sf_dtype cluster_shape_mn=(2, 1), sf_vec_size=sf_vec_size, sf_dtype=sf_dtype, - vector_f32=False, + vector_f32=vector_f32, discrete_col_sfd=False, act_func="dsituglu", request=request, From e8407f244991d9597b3f48a3010c4aec180fda48 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Wed, 19 Aug 2026 17:48:14 +0800 Subject: [PATCH 7/9] [CuTeDSL] Validate dense dSiTU-GLU backward Signed-off-by: Harry Zhou --- .../gemm_fusions/grouped_gemm_dglu.md | 4 ++- ...moe_blockscaled_grouped_gemm_dglu_dbias.py | 2 +- .../grouped_gemm/test_grouped_gemm_dglu.py | 7 ++-- .../test_grouped_gemm_dswiglu_utils.py | 36 ++++++++++++++----- 4 files changed, 34 insertions(+), 15 deletions(-) diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md index e3d38bffa..5468c90eb 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md @@ -147,7 +147,9 @@ $$ $$ and returns `ref * prob * T_u * dT_g/dG` and -`ref * prob * T_g * dT_u/dU`. `dprob` accumulates `ref * T_g * T_u`. +`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. diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py index 621bd6134..10f26af6d 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py @@ -1906,7 +1906,7 @@ def dsituglu( # The K3-default specialization always uses packed FP32x2 arithmetic. This is an # activation-specific optimization and is independent of the generic vectorized_f32 knob. - if cutlass.const_expr(self.situ_beta1 == 4.0): + if cutlass.const_expr(beta1 == 4.0): fmul2 = partial(cute.arch.mul_packed_f32x2, rnd="rn", ftz=False) fadd2 = partial(cute.arch.add_packed_f32x2, rnd="rn", ftz=False) square_alpha2 = (square_alpha, square_alpha) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py index cabccbb39..f3d383209 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py @@ -758,7 +758,6 @@ def _test_grouped_gemm_dglu_dense_wrapper( act_func="dswiglu", situ_beta1=4.0, situ_beta2=25.0, - skip_reference=False, ): try: from cudnn import grouped_gemm_dglu_wrapper_sm100 @@ -863,15 +862,14 @@ def _test_grouped_gemm_dglu_dense_wrapper( pytest.skip(f"Unsupported testcase: {e}") torch.cuda.synchronize() - if not skip_reference: - check_ref_grouped_gemm_dswiglu(inputs, wrapper_outputs, cfg, skip_ref=cfg["skip_ref"]) + check_ref_grouped_gemm_dswiglu(inputs, wrapper_outputs, cfg, skip_ref=cfg["skip_ref"]) return inputs, wrapper_outputs, cfg @pytest.mark.L0 @torch_fork_set_rng(seed=0) def test_grouped_gemm_dglu_dense_wrapper_dsituglu_mxfp8(request): - """Smoke-test the dense MXFP8 compile and execute path for dSiTU-GLU.""" + """Validate dense MXFP8 dSiTU-GLU outputs against the PyTorch reference.""" _test_grouped_gemm_dglu_dense_wrapper( ab_dtype=torch.float8_e4m3fn, @@ -888,7 +886,6 @@ def test_grouped_gemm_dglu_dense_wrapper_dsituglu_mxfp8(request): discrete_col_sfd=False, request=request, act_func="dsituglu", - skip_reference=True, ) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_utils.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_utils.py index f6630b7af..f84ebc167 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_utils.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_utils.py @@ -375,8 +375,11 @@ def run_grouped_gemm_dswiglu_ref( d_dtype: torch.dtype = torch.float32, sf_vec_size: int = 16, sf_dtype: torch.dtype = torch.float8_e8m0fnu, + act_func: str = "dswiglu", + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, ) -> Dict[str, torch.Tensor]: - """Run reference implementation for grouped GEMM dSwiGLU backward. + """Run reference implementation for grouped GEMM dGLU backward. Based on the reference in continugous_blockscaled_grouped_gemm_dswiglu_quant_fusion.py @@ -407,6 +410,9 @@ def run_grouped_gemm_dswiglu_ref( :param d_dtype: Output D tensor dtype :param sf_vec_size: Scale factor vector size :param sf_dtype: Scale factor dtype + :param act_func: Activation function (``"dswiglu"`` or ``"dsituglu"``) + :param situ_beta1: Gate tanh scale for dSiTU-GLU + :param situ_beta2: Up-branch tanh scale for dSiTU-GLU :return: Dictionary of reference tensors """ n, k, l = b_ref.shape @@ -454,19 +460,30 @@ def run_grouped_gemm_dswiglu_ref( c_input = c_full.index_select(dim=1, index=dest_idx_ab) # shape [M, N, L] c_gate = c_full.index_select(dim=1, index=dest_idx_glu) # shape [M, N, L] sig = torch.sigmoid(c_gate) - swish = c_gate * sig + if act_func == "dsituglu": + gate_tanh = torch.tanh(c_gate / situ_beta1) + up_tanh = torch.tanh(c_input / situ_beta2) + gate_value = situ_beta1 * gate_tanh * sig + up_value = situ_beta2 * up_tanh + gate_grad = (1.0 - gate_tanh.square()) * sig + gate_grad = gate_grad + situ_beta1 * gate_tanh * sig * (1.0 - sig) + up_grad = 1.0 - up_tanh.square() + ref_dprob = gate_value * up_value * ref + prob = prob_tensor.expand(-1, n, -1) + ab = ref * prob * gate_value * up_grad + dswiglu = ref * prob * up_value * gate_grad + else: + swish = c_gate * sig + ref_dprob = swish * c_input * ref + prob = prob_tensor.expand(-1, n, -1) + ab = ref * prob * swish + dswiglu = ref * prob * c_input * sig * (1 + c_gate * (1 - sig)) # Step 3: Compute dprob reference - ref_dprob = swish * c_input * ref chunk_sums = [torch.sum(chunk, dim=1, keepdim=True) for chunk in torch.split(ref_dprob, 32, dim=1)] ref_dprob = torch.sum(torch.cat(chunk_sums, dim=1), dim=1, keepdim=True) # (m, 1, l) ref_tensors["dprob_ref"] = ref_dprob - # Step 4: Compute dSwiGLU formulas - prob = prob_tensor.expand(-1, n, -1) - ab = ref * prob * swish - dswiglu = ref * prob * c_input * sig * (1 + c_gate * (1 - sig)) - # Step 5: Interleave [dswiglu, ab] back into swizzled [M, N, 1] by 32-wide blocks ref_d = torch.empty_like(c_full) # Place AB blocks at even positions @@ -561,6 +578,9 @@ def check_ref_grouped_gemm_dswiglu( d_dtype=cfg["d_dtype"], sf_vec_size=cfg["sf_vec_size"], sf_dtype=cfg["sf_dtype"], + act_func=cfg.get("act_func", "dswiglu"), + situ_beta1=cfg.get("situ_beta1", 4.0), + situ_beta2=cfg.get("situ_beta2", 25.0), ) torch.cuda.synchronize() From c6ee1633ea68bae4de55e6c23365e22030a4dd36 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Wed, 19 Aug 2026 23:38:09 +0800 Subject: [PATCH 8/9] [CuTeDSL] Cast dSiTU scaling factors to FP32 Signed-off-by: Harry Zhou --- ...moe_blockscaled_grouped_gemm_dglu_dbias.py | 17 +++++++----- .../grouped_gemm/test_grouped_gemm_dglu.py | 27 ++++++++++++++++--- 2 files changed, 34 insertions(+), 10 deletions(-) diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py index 10f26af6d..af6ea1aa0 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py @@ -1903,19 +1903,22 @@ def dsituglu( beta2_f32 = cutlass.Float32(beta2) beta1_rcp = cutlass.Float32(1.0 / beta1) beta2_rcp = cutlass.Float32(1.0 / beta2) + square_alpha_f32 = cutlass.Float32(square_alpha) + beta_val_f32 = cutlass.Float32(beta_val) + mprob_f32 = cutlass.Float32(mProb) # The K3-default specialization always uses packed FP32x2 arithmetic. This is an # activation-specific optimization and is independent of the generic vectorized_f32 knob. if cutlass.const_expr(beta1 == 4.0): fmul2 = partial(cute.arch.mul_packed_f32x2, rnd="rn", ftz=False) fadd2 = partial(cute.arch.add_packed_f32x2, rnd="rn", ftz=False) - square_alpha2 = (square_alpha, square_alpha) - beta_val2 = (beta_val, beta_val) + square_alpha2 = (square_alpha_f32, square_alpha_f32) + beta_val2 = (beta_val_f32, beta_val_f32) beta1_2 = (beta1_f32, beta1_f32) beta2_2 = (beta2_f32, beta2_f32) beta1_rcp2 = (beta1_rcp, beta1_rcp) beta2_rcp2 = (beta2_rcp, beta2_rcp) - mprob2 = (mProb, mProb) + mprob2 = (mprob_f32, mprob_f32) ones2 = (cutlass.Float32(1.0), cutlass.Float32(1.0)) halves2 = (cutlass.Float32(0.5), cutlass.Float32(0.5)) twos2 = (cutlass.Float32(2.0), cutlass.Float32(2.0)) @@ -1992,9 +1995,9 @@ def dsituglu( return dgate_vec.load(), dup_vec.load(), dprob_swiglu for i in cutlass.range_constexpr(cute.size(acc_vec)): - grad = acc_vec[i] * square_alpha - gate = gate_vec[i].to(cutlass.Float32) * beta_val - up = up_vec[i].to(cutlass.Float32) * beta_val + grad = acc_vec[i] * square_alpha_f32 + gate = gate_vec[i].to(cutlass.Float32) * beta_val_f32 + up = up_vec[i].to(cutlass.Float32) * beta_val_f32 gate_tanh = cute.math.tanh(gate * beta1_rcp, fastmath=True) up_tanh = cute.math.tanh(up * beta2_rcp, fastmath=True) sigmoid = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True)) @@ -2005,7 +2008,7 @@ def dsituglu( up_grad = cutlass.Float32(1.0) - up_tanh * up_tanh activation_grad = grad if cutlass.const_expr(self.has_prob): - activation_grad = grad * mProb + activation_grad = grad * mprob_f32 dgate_vec[i] = activation_grad * up_value * gate_grad dup_vec[i] = activation_grad * gate_value * up_grad if cutlass.const_expr(self.generate_dprob): diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py index f3d383209..985eafb4c 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py @@ -868,8 +868,27 @@ def _test_grouped_gemm_dglu_dense_wrapper( @pytest.mark.L0 @torch_fork_set_rng(seed=0) -def test_grouped_gemm_dglu_dense_wrapper_dsituglu_mxfp8(request): - """Validate dense MXFP8 dSiTU-GLU outputs against the PyTorch reference.""" +@pytest.mark.parametrize( + "alpha_beta_dtype,situ_beta1,vector_f32", + [ + pytest.param(torch.float32, 4.0, False, id="fp32-default-scalar"), + pytest.param(torch.bfloat16, 4.0, False, id="bf16-default-scalar"), + pytest.param(torch.bfloat16, 4.0, True, id="bf16-default-vector"), + pytest.param(torch.bfloat16, 2.0, False, id="bf16-generic-scalar"), + pytest.param(torch.bfloat16, 2.0, True, id="bf16-generic-vector"), + ], +) +def test_grouped_gemm_dglu_dense_wrapper_dsituglu_mxfp8(request, alpha_beta_dtype, situ_beta1, vector_f32): + """Validate dense MXFP8 dSiTU-GLU with caller-provided scaling dtypes. + + Transformer Engine supplies BF16 alpha and beta tensors in its MXFP8 fused + grouped-MLP path. Cover both the K3-default packed specialization and the + generic-beta implementation, with the generic vectorization knob on and off. + """ + + def cast_scaling_tensors(inputs, _cfg): + inputs["alpha_tensor"] = inputs["alpha_tensor"].to(alpha_beta_dtype) + inputs["beta_tensor"] = inputs["beta_tensor"].to(alpha_beta_dtype) _test_grouped_gemm_dglu_dense_wrapper( ab_dtype=torch.float8_e4m3fn, @@ -882,10 +901,12 @@ def test_grouped_gemm_dglu_dense_wrapper_dsituglu_mxfp8(request): cluster_shape_mn=(2, 1), sf_vec_size=32, sf_dtype=torch.float8_e8m0fnu, - vector_f32=False, + vector_f32=vector_f32, discrete_col_sfd=False, request=request, + input_mutator=cast_scaling_tensors, act_func="dsituglu", + situ_beta1=situ_beta1, ) From 354d341911e2415e9b879f5175af9b57a1d6eb2c Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Thu, 20 Aug 2026 00:30:01 +0800 Subject: [PATCH 9/9] [CuTeDSL] Fuse SiTU-GLU with NVFP4 Hadamard Signed-off-by: Harry Zhou --- .../gemm_fusions/grouped_gemm_glu_hadamard.md | 69 +++++++---- .../gemm/cutedsl/grouped/glu_hadamard/api.py | 58 ++++++++- ...e_blockscaled_grouped_gemm_glu_hadamard.py | 37 +++++- .../test_grouped_gemm_glu_hadamard.py | 114 +++++++++++++++++- 4 files changed, 249 insertions(+), 29 deletions(-) diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md index 2166ad4e0..477cee760 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md +++ b/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,6 +161,7 @@ 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), @@ -158,7 +169,9 @@ op = GroupedGemmGluHadamardSm100( 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)`. --- diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/api.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/api.py index c1b973238..d4d43bb94 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/api.py @@ -6,6 +6,7 @@ from __future__ import annotations import logging +import math import os from typing import Any, Optional, Tuple @@ -81,6 +82,8 @@ def __init__( vector_f32: bool = False, m_aligned: int = 256, act_func: str = "swiglu", + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, use_dynamic_sched: bool = False, use_tmem_post_rht_amax: bool = False, ): @@ -157,6 +160,17 @@ def __init__( self.vector_f32 = vector_f32 self.m_aligned = m_aligned self.act_func = act_func + self.situ_beta1 = float(situ_beta1) + self.situ_beta2 = float(situ_beta2) + if self.act_func == "situglu": + self._value_error_if( + not math.isfinite(self.situ_beta1) or self.situ_beta1 <= 0.0, + f"situ_beta1 must be finite and positive, got {self.situ_beta1}", + ) + self._value_error_if( + not math.isfinite(self.situ_beta2) or self.situ_beta2 <= 0.0, + f"situ_beta2 must be finite and positive, got {self.situ_beta2}", + ) self.use_dynamic_sched = use_dynamic_sched self.use_tmem_post_rht_amax = use_tmem_post_rht_amax self._kernel = BlockScaledMoEGroupedGemmGluHadamardKernel @@ -258,8 +272,8 @@ def check_support(self) -> bool: self._value_error_if(self.sf_vec_size not in [16, 32], f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") self._value_error_if( - self.act_func not in ["swiglu", "geglu", "srelu"], - f"act_func must be 'swiglu', 'geglu', or 'srelu', got {self.act_func}", + self.act_func not in ["swiglu", "geglu", "situglu", "srelu"], + f"act_func must be 'swiglu', 'geglu', 'situglu', or 'srelu', got {self.act_func}", ) self._value_error_if( not self.use_2cta_instrs or self.mma_tiler_mn != (256, 256), f"Hadamard fusion requires mma_tiler_mn=(256, 256), got {self.mma_tiler_mn}" @@ -333,6 +347,7 @@ def compile(self) -> None: weight_mode=self.weight_mode, use_dynamic_sched=self.use_dynamic_sched, act_func=self.act_func, + situ_beta1=self.situ_beta1, enable_bias=self.bias_desc is not None, use_tmem_post_rht_amax=self.use_tmem_post_rht_amax, ) @@ -442,6 +457,8 @@ def compile(self) -> None: max_active_clusters, fake_stream, cached_linear_offset, + cutlass.Float32(self.situ_beta1), + cutlass.Float32(self.situ_beta2), options="--enable-tvm-ffi", ) @@ -462,6 +479,8 @@ def tensor_api( post_rht_amax_tensor: Optional[torch.Tensor], bias_tensor: Optional[torch.Tensor], stream: cuda.CUstream, + situ_beta1: float, + situ_beta2: float, ) -> None: compiled_kernel( a_tensor, @@ -483,6 +502,8 @@ def tensor_api( bias_tensor, stream, cached_linear_offset, + cutlass.Float32(situ_beta1), + cutlass.Float32(situ_beta2), ) self._compiled_kernel = tensor_api @@ -506,6 +527,8 @@ def tensor_api( post_rht_amax_tensor: Optional[torch.Tensor], bias_tensor: Optional[torch.Tensor], stream: cuda.CUstream, + situ_beta1: float, + situ_beta2: float, ) -> None: compiled_kernel( a_tensor, @@ -527,6 +550,8 @@ def tensor_api( bias_tensor, stream, cached_linear_offset, + cutlass.Float32(situ_beta1), + cutlass.Float32(situ_beta2), ) self._compiled_kernel = tensor_api @@ -548,6 +573,8 @@ def execute( amax_tensor: Optional[torch.Tensor] = None, post_rht_amax_tensor: Optional[torch.Tensor] = None, bias_tensor: Optional[torch.Tensor] = None, + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, current_stream: Optional[cuda.CUstream] = None, ) -> None: import torch @@ -559,6 +586,19 @@ def execute( return if current_stream is None: current_stream = cuda.CUstream(torch.cuda.current_stream(a_tensor.device).cuda_stream) + if self.act_func == "situglu": + self._value_error_if( + not math.isfinite(situ_beta1) or situ_beta1 <= 0.0, + f"situ_beta1 must be finite and positive, got {situ_beta1}", + ) + self._value_error_if( + not math.isfinite(situ_beta2) or situ_beta2 <= 0.0, + f"situ_beta2 must be finite and positive, got {situ_beta2}", + ) + self._value_error_if( + float(situ_beta1) != self.situ_beta1, + "situ_beta1 is specialized at compile time; construct and compile " f"the API with situ_beta1={situ_beta1}", + ) if hadamard_tensor is None: hadamard_tensor = self.hadamard_tensor else: @@ -586,6 +626,8 @@ def execute( post_rht_amax_tensor, bias_tensor, current_stream, + situ_beta1, + situ_beta2, ) else: if b_ptrs is None or sfb_ptrs is None: @@ -605,6 +647,8 @@ def execute( post_rht_amax_tensor, bias_tensor, current_stream, + situ_beta1, + situ_beta2, ) @@ -636,6 +680,8 @@ def grouped_gemm_glu_hadamard_wrapper_sm100( vector_f32: bool = False, m_aligned: int = 256, act_func: str = "swiglu", + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, use_dynamic_sched: bool = False, use_tmem_post_rht_amax: bool = False, current_stream: Optional[cuda.CUstream] = None, @@ -712,9 +758,11 @@ def dynamic_m_tensor_signature( stride_signature = tuple(None if idx in dynamic_stride_dims else value for idx, value in enumerate(tensor.stride())) return static_shape_suffix, stride_signature, tensor.dtype + situ_beta1_cache_signature = float(situ_beta1) if act_func == "situglu" else None cache_key = ( weight_mode, act_func, + situ_beta1_cache_signature, a_tensor.shape[1:], tuple(b_tensor.shape) if is_dense else b_shape, c_tensor.shape[1:], @@ -763,6 +811,8 @@ def dynamic_m_tensor_signature( vector_f32=vector_f32, m_aligned=m_aligned, act_func=act_func, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, use_dynamic_sched=use_dynamic_sched, use_tmem_post_rht_amax=use_tmem_post_rht_amax, ) @@ -794,6 +844,8 @@ def dynamic_m_tensor_signature( amax_tensor=amax_tensor, post_rht_amax_tensor=post_rht_amax_tensor, bias_tensor=bias_tensor, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, current_stream=current_stream, ) else: @@ -810,6 +862,8 @@ def dynamic_m_tensor_signature( amax_tensor=amax_tensor, post_rht_amax_tensor=post_rht_amax_tensor, bias_tensor=bias_tensor, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, current_stream=current_stream, ) return TupleDict(c_tensor=c_tensor, d_tensor=d_tensor, amax_tensor=amax_tensor, post_rht_amax_tensor=post_rht_amax_tensor) diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py index ddab650f7..c7532fd43 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py @@ -145,6 +145,7 @@ def __init__( weight_mode: MoEWeightMode = MoEWeightMode.DISCRETE, use_dynamic_sched: bool = False, act_func: str = "swiglu", + situ_beta1: float = 4.0, enable_bias: bool = False, use_tmem_post_rht_amax: bool = False, ): @@ -234,7 +235,8 @@ def __init__( self.num_epilog_warps = len(self.epilog_act_warp_id) # = 4 self.act_func = act_func - if act_func not in ["swiglu", "geglu", "srelu"]: + self.situ_beta1 = float(situ_beta1) + if act_func not in ["swiglu", "geglu", "situglu", "srelu"]: raise ValueError(f"Invalid activation function: {act_func}") def _setup_attributes(self): @@ -552,12 +554,21 @@ def __call__( stream: cuda.CUstream, epilogue_op: cutlass.Constexpr = lambda x: x, linear_offset: cutlass.Float32 = 0.0, + situ_beta1: cutlass.Float32 = 4.0, + situ_beta2: cutlass.Float32 = 25.0, ): """Execute the MoE GEMM + GLU + Hadamard kernel. Dense mode: ``b`` and ``sfb`` are 3-D cute.Tensor (N, K, L). Discrete mode: ``b`` and ``sfb`` are cute.Pointer to device int64[] arrays of per-expert base addresses. + + ``situ_beta1`` and ``situ_beta2`` configure SiTU-GLU: + + out = beta1 * tanh(gate / beta1) * sigmoid(gate) + * beta2 * tanh(up / beta2) + + They are ignored unless ``act_func == "situglu"``. """ self.a_dtype: Type[cutlass.Numeric] = a.element_type self.b_dtype: Type[cutlass.Numeric] = a.element_type @@ -865,6 +876,8 @@ class SharedStorage: self.sched_params, epilogue_op, linear_offset, + situ_beta1, + situ_beta2, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], @@ -1093,6 +1106,23 @@ def swiglu_act(self, tCompute, acc_vec_up, acc_vec_gate, mProb): tCompute[i] = acc_vec_up[i] * silu_f32(acc_vec_gate[i], fastmath=True) tCompute[i] = tCompute[i] * mProb + @cute.jit + def situglu_act(self, tCompute, acc_vec_up, acc_vec_gate, mProb, beta1, beta2): + beta1_rcp = cutlass.Float32(1.0 / self.situ_beta1) + beta2_rcp = cute.arch.rcp_approx(beta2) + beta_product = beta1 * beta2 + for i in cutlass.range_constexpr(cute.size(tCompute)): + gate = acc_vec_gate[i] + up = acc_vec_up[i] + gate_tanh = cute.math.tanh(gate * beta1_rcp, fastmath=True) + up_tanh = cute.math.tanh(up * beta2_rcp, fastmath=True) + if cutlass.const_expr(self.situ_beta1 == 4.0): + # For a = tanh(gate / 4), sigmoid(gate) = 1/2 + a / (1 + a^2). + sigmoid = cutlass.Float32(0.5) + gate_tanh * cute.arch.rcp_approx(cutlass.Float32(1.0) + gate_tanh * gate_tanh) + else: + sigmoid = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True)) + tCompute[i] = beta_product * gate_tanh * sigmoid * up_tanh * mProb + @cute.jit def srelu_act(self, tCompute, acc_vec, mProb): acc_relu = cute.where(acc_vec > 0, acc_vec, cute.full_like(acc_vec, 0)) @@ -1265,6 +1295,8 @@ def kernel( sched_params: MoESchedulerParams, epilogue_op: cutlass.Constexpr, linear_offset: cutlass.Float32 = 0.0, + situ_beta1: cutlass.Float32 = 4.0, + situ_beta2: cutlass.Float32 = 25.0, ): """GPU device kernel: MoE persistent GEMM + GLU + Hadamard (pingpong epilogue).""" warp_idx = cute.arch.warp_idx() @@ -2208,6 +2240,9 @@ def kernel( elif cutlass.const_expr(self.act_func == "swiglu"): acc_vec_up = tTR_rAcc_up.load() self.swiglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb) + elif cutlass.const_expr(self.act_func == "situglu"): + acc_vec_up = tTR_rAcc_up.load() + self.situglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb, situ_beta1, situ_beta2) if cutlass.const_expr(self.generate_amax): thread_tile_amax = self.amax_reduction_per_thread(tCompute, thread_tile_amax) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard.py index 5a2476aeb..502ba80d3 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard.py @@ -55,7 +55,7 @@ def _apply_transpose_hadamard(d_ref: torch.Tensor, aligned_group_m_list) -> torc return out.unsqueeze(-1) -def _run_grouped_gemm_glu_ref(inputs: Dict, act_func: str) -> Dict: +def _run_grouped_gemm_glu_ref(inputs: Dict, act_func: str, situ_beta1: float = 4.0, situ_beta2: float = 25.0) -> Dict: n, _, l = inputs["b_ref"].shape n_out = n // 2 valid_m = inputs["valid_m"] @@ -102,6 +102,8 @@ def _run_grouped_gemm_glu_ref(inputs: Dict, act_func: str) -> Dict: ref_gate = torch.clamp(ref_gate, max=7.0) ref_up = torch.clamp(ref_up, min=-7.0, max=7.0) ref_after_glu = (ref_up + 1.0) * ref_gate * torch.sigmoid(1.702 * ref_gate) + elif act_func == "situglu": + ref_after_glu = situ_beta1 * torch.tanh(ref_gate / situ_beta1) * torch.sigmoid(ref_gate) * situ_beta2 * torch.tanh(ref_up / situ_beta2) else: raise ValueError(f"Unsupported act_func {act_func}") @@ -109,8 +111,16 @@ def _run_grouped_gemm_glu_ref(inputs: Dict, act_func: str) -> Dict: return {"c_ref": ref.clone(), "d_ref": ref_after_glu} -def _check_reference(inputs: Dict, outputs: Dict, cfg: Dict, *, act_func: str) -> None: - ref_tensors = _run_grouped_gemm_glu_ref(inputs, act_func) +def _check_reference( + inputs: Dict, + outputs: Dict, + cfg: Dict, + *, + act_func: str, + situ_beta1: float = 4.0, + situ_beta2: float = 25.0, +) -> None: + ref_tensors = _run_grouped_gemm_glu_ref(inputs, act_func, situ_beta1, situ_beta2) torch.testing.assert_close( outputs["c_tensor"][: inputs["valid_m"]].cpu().float(), @@ -230,7 +240,17 @@ def _run_compile_execute(request, *, ab_dtype, sf_dtype, sf_vec_size, act_func=" _check_reference(inputs, outputs, cfg, act_func=act_func) -def _run_wrapper(request, *, ab_dtype, sf_dtype, sf_vec_size, act_func="swiglu", enable_bias=False): +def _run_wrapper( + request, + *, + ab_dtype, + sf_dtype, + sf_vec_size, + act_func="swiglu", + enable_bias=False, + situ_beta1=4.0, + situ_beta2=25.0, +): cfg = _make_cfg(request, ab_dtype=ab_dtype, sf_dtype=sf_dtype, sf_vec_size=sf_vec_size, enable_bias=enable_bias) inputs = allocate_grouped_gemm_input_tensors( n=cfg["n"], @@ -266,9 +286,11 @@ def _run_wrapper(request, *, ab_dtype, sf_dtype, sf_vec_size, act_func="swiglu", vector_f32=cfg["vector_f32"], m_aligned=cfg["m_aligned"], act_func=act_func, + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, ) - _check_reference(inputs, outputs, cfg, act_func=act_func) + _check_reference(inputs, outputs, cfg, act_func=act_func, situ_beta1=situ_beta1, situ_beta2=situ_beta2) def _run_discrete_wrapper(request, *, ab_dtype, sf_dtype, sf_vec_size, act_func="swiglu"): @@ -342,6 +364,21 @@ def test_grouped_gemm_glu_hadamard_wrapper_fp4(request, ab_dtype, sf_dtype, sf_v ) +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize(("situ_beta1", "situ_beta2"), [(4.0, 25.0), (2.0, 8.0)]) +def test_grouped_gemm_glu_hadamard_wrapper_situglu_fp4(request, situ_beta1, situ_beta2): + _run_wrapper( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e8m0fnu, + sf_vec_size=16, + act_func="situglu", + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, + ) + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) def test_grouped_gemm_glu_hadamard_wrapper_discrete_fp4(request): @@ -431,3 +468,70 @@ def counted_compile(self): assert compile_count["value"] == 1 assert len(grouped_gemm_glu_hadamard_api._cache_of_GroupedGemmGluHadamardSm100Objects) == 1 grouped_gemm_glu_hadamard_api._cache_of_GroupedGemmGluHadamardSm100Objects.clear() + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_hadamard_wrapper_cache_specializes_situ_beta1_only(request, monkeypatch): + from cudnn import grouped_gemm_glu_hadamard_wrapper_sm100 + from cudnn.gemm.cutedsl.grouped.glu_hadamard import api as grouped_gemm_glu_hadamard_api + + grouped_gemm_glu_hadamard_api._cache_of_GroupedGemmGluHadamardSm100Objects.clear() + + compile_count = {"value": 0} + + def counted_compile(self): + compile_count["value"] += 1 + return None + + monkeypatch.setattr(grouped_gemm_glu_hadamard_api.GroupedGemmGluHadamardSm100, "compile", counted_compile) + monkeypatch.setattr(grouped_gemm_glu_hadamard_api.GroupedGemmGluHadamardSm100, "check_support", lambda self: True) + monkeypatch.setattr(grouped_gemm_glu_hadamard_api.GroupedGemmGluHadamardSm100, "execute", lambda self, **kwargs: None) + + cfg = _make_cfg( + request, + ab_dtype=torch.float4_e2m1fn_x2, + sf_dtype=torch.float8_e8m0fnu, + sf_vec_size=16, + ) + inputs = allocate_grouped_gemm_input_tensors( + n=cfg["n"], + k=cfg["k"], + l=cfg["l"], + group_m_list=cfg["group_m_list"], + ab_dtype=cfg["ab_dtype"], + sf_dtype=cfg["sf_dtype"], + sf_vec_size=cfg["sf_vec_size"], + m_aligned=cfg["m_aligned"], + b_major=cfg["b_major"], + enable_bias=False, + ) + + for situ_beta1, situ_beta2 in ((4.0, 25.0), (4.0, 8.0), (2.0, 25.0), (4.0, 25.0)): + grouped_gemm_glu_hadamard_wrapper_sm100( + a_tensor=inputs["a_tensor"], + b_tensor=inputs["b_tensor"], + sfa_tensor=inputs["sfa_tensor"], + sfb_tensor=inputs["sfb_tensor"], + padded_offsets=inputs["padded_offsets_tensor"], + alpha_tensor=inputs["alpha_tensor"], + prob_tensor=inputs["prob_tensor"], + acc_dtype=cfg["acc_dtype"], + c_dtype=cfg["c_dtype"], + d_dtype=cfg["d_dtype"], + cd_major=cfg["cd_major"], + mma_tiler_mn=cfg["mma_tiler_mn"], + cluster_shape_mn=cfg["cluster_shape_mn"], + sf_vec_size=cfg["sf_vec_size"], + vector_f32=cfg["vector_f32"], + m_aligned=cfg["m_aligned"], + act_func="situglu", + situ_beta1=situ_beta1, + situ_beta2=situ_beta2, + ) + + # beta1 changes the generated activation code and therefore creates a new + # compiled API. beta2 is a runtime FP32 scalar and reuses the beta1=4 API. + assert compile_count["value"] == 2 + assert len(grouped_gemm_glu_hadamard_api._cache_of_GroupedGemmGluHadamardSm100Objects) == 2 + grouped_gemm_glu_hadamard_api._cache_of_GroupedGemmGluHadamardSm100Objects.clear()