Grouped GEMM + dGLU (SM100)

This is an experimental API and subject to change.

JAX support

Supports JAX arrays on the BF16 backend in discrete weight mode (dswiglu and dgeglu), including generate_dbias=True and caller-provided zero-initialized dprob. Dense b_tensor and the block-scaled backend (MMA-interleaved scale-factor layouts) are not expressible as JAX arrays and raise clear errors. The wrapper is eager, on the CUDA legacy default stream: block_until_ready inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes.

For jitted JAX programs use the jax.jit-compatible XLA custom-call entry point grouped_gemm_dglu_jax_sm100 (built on cudnn.jax.call; discrete mode): dprob and (with generate_dbias=True) dbias come back as bridge-managed zero-initialized accumulator outputs — no caller-zeroed buffers, no manual synchronization. Under tracing the padded_offsets values cannot be host-validated, and the per-expert weight buffers behind b_ptrs must stay alive and unmoved across every execution of the traced computation.

Overview

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, 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:

• Dense mode: All expert weights packed into a single contiguous (N, K, L) tensor

• Discrete mode: Per-expert weight pointers (no weight stacking required)

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).

Backend dispatch

Operand contract	Selected backend
A and B are BF16	BF16
matching supported FP4/FP8 A and B plus scale descriptors	block-scaled

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

Pass sfa_tensor=None, sfb_tensor=None (or sfb_ptrs=None), and norm_const_tensor=None; keep sf_vec_size=16, discrete_col_sfd=False, and epilogue_op=None. BF16 also uses the source GeGLU constants geglu_alpha=1.702, glu_clamp_max=7, and glu_clamp_min=-7.

Tensors, layouts, and equation

For padded rows M, reduction dimension K, compact gradient width N, and L experts, BF16 uses:

• A: (M, K, 1), stride (K, 1, M*K), BF16;

• dense B: (N, K, L), K-major stride (K, 1, N*K), BF16;

• discrete b_ptrs: contiguous CUDA int64 pointers to expert (N, K) BF16 matrices, with n=N, b_dtype=torch.bfloat16, and b_major="k" or "n";

• C: (M, 2N, 1), stride (2N, 1, 2M*N), BF16/FP16/FP32;

• padded_offsets: (L,) int32 cumulative 256-aligned ends;

• alpha, beta: (L,) FP32; prob: (M, 1, 1) FP32;

• caller-zeroed dprob: (M, 1, 1), stride (1, 1, 1), FP32;

• D_row: (M, 2N, 1), stride (2N, 1, 2M*N), BF16/FP16/FP32;

• caller-zeroed optional dbias: (L, 2N, 1), stride (2N, 1, 1), BF16.

For expert g, the compact GEMM gradient and scaled forward activation are

\( R_g = \alpha_g^2 A_g B_g^T, \qquad X_g = \beta_g C_g. \)

Split X into alternating 32-wide gate/input blocks. For dSwiGLU, with s = sigmoid(gate):

\( d\mathrm{input} = R\,\mathrm{prob}\,(\mathrm{gate}\,s), \)

\( d\mathrm{gate} = R\,\mathrm{prob}\,\mathrm{input}\,s (1 + \mathrm{gate}(1-s)). \)

For dGeGLU, distinguish raw values, clamped activation values, and the source’s value-bearing filters:

raw_gate = gate(X)
raw_input = input(X)
clamped_gate = min(raw_gate, 7)
clamped_input = clamp(raw_input, -7, 7)
gate_filter = raw_gate if raw_gate <= 7 else 0
input_filter = raw_input if -7 <= raw_input <= 7 else 0
s = sigmoid(1.702 * clamped_gate)

With the default linear_offset=1, the kernel computes

\( d\mathrm{gate} = R\,\mathrm{prob}\, (\mathrm{clamped\_input}+\mathrm{linear\_offset})\,s (1 + 1.702\,\mathrm{clamped\_gate}(1-s))\, \mathrm{gate\_filter}, \)

\( d\mathrm{input} = R\,\mathrm{prob}\,\mathrm{clamped\_gate}\,s\, \mathrm{input\_filter}. \)

dprob accumulates the row sum of the matching unscaled activation times R; dbias is the per-expert row reduction of interleaved D_row. The pointer-array tensor is stream-recorded, while every pointed allocation must remain alive and unchanged until that stream completes.

The wrapper return order is exactly d_row_tensor, d_col_tensor, dprob_tensor, dbias_tensor, amax_tensor, sfd_row_tensor, sfd_col_tensor. On BF16, d_col_tensor, amax_tensor, sfd_row_tensor, and sfd_col_tensor are None; dbias_tensor is None unless generate_dbias=True.

Block-scaled contract

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. dGLU backward epilogue: Fused backward computation using the forward C tensor (input/gate interleaved)

3. Optional quantized output: Produces row and column scale factors for downstream quantization

Shapes

Equations

For dSiTU-GLU, define

\( T_g=\beta_1\tanh(G/\beta_1)\sigma(G),\qquad T_u=\beta_2\tanh(U/\beta_2). \)

The fused backward computes

\( \frac{\partial T_g}{\partial G}= (1-\tanh^2(G/\beta_1))\sigma(G)+ \beta_1\tanh(G/\beta_1)\sigma(G)(1-\sigma(G)), \)

\( \frac{\partial T_u}{\partial U}=1-\tanh^2(U/\beta_2), \)

and returns ref * prob * T_u * dT_g/dG and ref * prob * T_g * dT_u/dU. dprob accumulates the reduction of ref * T_g * T_u across the output columns in 32-column chunks, producing shape (valid_m, 1, 1). The beta values are compile-time specialization values and therefore belong to the dGLU compiled-kernel cache key.

• Inputs

  • A: contiguous activation tensor across all groups, shape (valid_m, K, 1)

  • B (dense): weight tensor across all groups, shape (N, K, L)

  • B (discrete): per-expert weight pointers, b_ptrs shape (num_experts,) of int64

  • C: forward intermediate tensor with interleaved input/gate blocks, shape (valid_m, 2N, 1)

  • SFA: scale factor tensor for A, shape (32, 4, ceil(valid_m/128), 4, ceil(ceil(K/sf_vec_size)/4), 1)

  • SFB (dense): scale factor tensor for B, shape (32, 4, ceil(N/128), 4, ceil(ceil(K/sf_vec_size)/4), L)

  • SFB (discrete): per-expert SFB pointers, sfb_ptrs shape (num_experts,) of int64

  • padded_offsets: cumulative sum of aligned group M sizes, shape (L,). valid_m = padded_offsets[-1]

  • alpha: per-group scaling factors for GEMM, shape (L,)

  • beta: per-group scaling factors for C, shape (L,)

  • prob: per-row gating probabilities, shape (valid_m, 1, 1)

  • norm_const: normalization constant for FP8 quantization, shape (1,)

• Outputs

  • D_row: row-quantized dGLU output, shape (valid_m, 2N, 1)

  • D_col: column-quantized dGLU output, shape (valid_m, 2N, 1)

  • dprob: gradient of prob, shape (valid_m, 1, 1). Must be zero-initialized.

  • dbias (optional): per-expert bias gradient tensor, shape (L, 2N, 1)

  • SFD_row: row scale factors (when d_dtype is FP8), shape (32, 4, ceil(valid_m/128), 4, ceil(ceil((2N)/sf_vec_size)/4), 1)

  • SFD_col: column scale factors (when d_dtype is FP8), shape (32, 4, ceil((2N)/128), 4, ceil(ceil(valid_m/sf_vec_size)/4), 1)

  • amax: per-group amax (when d_dtype is bf16/float16), shape (L, 2, 1)

Step 1: Block-scaled grouped GEMM (per group g with rows m in [padded_offsets[g-1], padded_offsets[g])):

\( \text{ref}[m, n] = \alpha_g^2 \sum_{k} \text{dequantize}(A[m, k], \text{SFA}) \cdot \text{dequantize}(B[n, k, g], \text{SFB}) \)

Step 2: dGLU backward epilogue (performed with 32-column interleaving along 2N):

• Scale C by beta_g per group and deinterleave into input/gate halves by 32-wide blocks.

For dSwiGLU (act_func="dswiglu"):

• swish = gate * sigmoid(gate)

• dprob += sum(swish * input * ref) over 32-column chunks

• ab = ref * prob * swish

• dswiglu = ref * prob * input * sigmoid(gate) * (1 + gate * (1 - sigmoid(gate)))

• Interleave [ab, dswiglu] back into D_row/D_col with 32-column blocks.

For dGeGLU (act_func="dgeglu"): Uses sigmoid(1.702 * gate) scaling in the backward computation.

Step 3: Optional output quantization (when SFD outputs are generated):

\( \text{SFD_row}[m, n] = \text{norm_const} \cdot \max_{k \in \text{block}} |D[m, k]| \cdot \text{rcp_max} \)

\( D_{\text{quantized}}[m, n] = D[m, n] \cdot \frac{\text{norm_const}}{\text{SFD_row}[m, n]} \)

Diagram

 A (valid_m×K×1)    B (N×K×L) or b_ptrs    padded_offsets
 SFA                SFB or sfb_ptrs              |
   |                 |                           |
   |    +------------+                           |
   |    |                                        |
   v    v                                        v
  Dequantize → Grouped GEMM (per group ranges) → ref
                    |
                    | × alpha[group_idx]
                    v
               ref (valid_m×N×1)
                    |
 C (valid_m×2N×1) --× beta[group_idx]--> deinterleave 32-col blocks
                    |                    |
                    |                swish, sigmoid
                    |                    |
                    +--> dprob (sum over blocks)
                    |
                    +--> ab, dswiglu → interleave → D (valid_m×2N×1)
                                      |
                         +-----------+-----------+
                         |                       |
                         v                       v
                    Row Quantize            Col Quantize
                         |                       |
                         v                       v
                    D_row, SFD_row         D_col, SFD_col

---

API usage

BF16

High-level wrapper

import cudnn
import torch

dprob.zero_()
out = cudnn.grouped_gemm_dglu_wrapper_sm100(
    a_tensor=a,
    c_tensor=c,
    sfa_tensor=None,
    padded_offsets=padded_offsets,
    alpha_tensor=alpha,
    beta_tensor=beta,
    prob_tensor=prob,
    dprob_tensor=dprob,
    b_tensor=b,
    sfb_tensor=None,
    act_func="dswiglu",
    generate_dbias=True,
    use_dynamic_sched=True,
)
d_row, d_col, dprob, dbias, amax, sfd_row, sfd_col = out

# Discrete mode replaces the dense weight arguments.
out = cudnn.grouped_gemm_dglu_wrapper_sm100(
    a_tensor=a, c_tensor=c, sfa_tensor=None,
    padded_offsets=padded_offsets, alpha_tensor=alpha, beta_tensor=beta,
    prob_tensor=prob, dprob_tensor=dprob,
    b_ptrs=b_ptrs, sfb_ptrs=None, n=N, b_dtype=torch.bfloat16,
    act_func="dgeglu",
)

Class API

op = cudnn.GroupedGemmDgluSm100(
    sample_a=a, sample_c=c, sample_d_row=d_row, sample_d_col=None,
    sample_sfa=None, sample_padded_offsets=padded_offsets,
    sample_alpha=alpha, sample_beta=beta, sample_prob=prob,
    sample_dprob=dprob, sample_dbias=dbias,
    sample_b=b, sample_sfb=None, act_func="dswiglu",
)
assert op.check_support()
op.compile()
dprob.zero_()
dbias.zero_()
op.execute(
    a_tensor=a, c_tensor=c, d_row_tensor=d_row, d_col_tensor=None,
    sfa_tensor=None, padded_offsets=padded_offsets, alpha_tensor=alpha,
    beta_tensor=beta, prob_tensor=prob, dprob_tensor=dprob,
    dbias_tensor=dbias, b_tensor=b, sfb_tensor=None,
)

use_dynamic_sched=False selects static scheduling; True caches a dynamic-M callable for compatible shapes. Cache keys retain compile-sensitive layouts, dtypes, activation, dbias policy, scheduler, tiles/clusters, features, and overlap margin.

Block-scaled

High-level wrapper

Dense mode:

from cudnn import grouped_gemm_dglu_wrapper_sm100
from cuda.bindings import driver as cuda

stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

outputs = grouped_gemm_dglu_wrapper_sm100(
    a_tensor=a,
    c_tensor=c,
    sfa_tensor=sfa,
    padded_offsets=padded_offsets,
    alpha_tensor=alpha,
    beta_tensor=beta,
    prob_tensor=prob,
    dprob_tensor=dprob,
    generate_dbias=True,
    # Dense mode weights:
    b_tensor=b,
    sfb_tensor=sfb,
    # Common:
    norm_const_tensor=norm_const,
    acc_dtype=torch.float32,
    d_dtype=torch.bfloat16,
    cd_major="n",
    mma_tiler_mn=(256, 256),
    cluster_shape_mn=(2, 1),
    sf_vec_size=32,
    act_func="dswiglu",
    epilogue_op=None,
    current_stream=stream,
)

# dictionary access:
d_row = outputs["d_row_tensor"]
d_col = outputs["d_col_tensor"]
dprob = outputs["dprob_tensor"]
dbias = outputs["dbias_tensor"]
amax = outputs["amax_tensor"]
sfd_row = outputs["sfd_row_tensor"]
sfd_col = outputs["sfd_col_tensor"]

# or tuple unpacking:
d_row, d_col, dprob, dbias, amax, sfd_row, sfd_col = outputs

Discrete mode:

outputs = grouped_gemm_dglu_wrapper_sm100(
    a_tensor=a,
    c_tensor=c,
    sfa_tensor=sfa,
    padded_offsets=padded_offsets,
    alpha_tensor=alpha,
    beta_tensor=beta,
    prob_tensor=prob,
    dprob_tensor=dprob,
    # Discrete mode weights:
    b_ptrs=b_ptrs,
    sfb_ptrs=sfb_ptrs,
    n=n_dim,
    b_dtype=torch.uint8,
    b_major="k",
    # Common:
    act_func="dgeglu",
    current_stream=stream,
)

Class API

Dense mode:

from cudnn import GroupedGemmDgluSm100
from cuda.bindings import driver as cuda

stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

api = GroupedGemmDgluSm100(
    sample_a=a,
    sample_c=c,
    sample_d_row=d_row,
    sample_d_col=d_col,
    sample_sfa=sfa,
    sample_padded_offsets=padded_offsets,
    sample_alpha=alpha,
    sample_beta=beta,
    sample_prob=prob,
    sample_dprob=dprob,
    sample_dbias=dbias,
    # Dense mode:
    sample_b=b,
    sample_sfb=sfb,
    # Optional quantization outputs
    sample_sfd_row=sfd_row,
    sample_sfd_col=sfd_col,
    sample_amax=amax,
    sample_norm_const=norm_const,
    # Configuration
    acc_dtype=torch.float32,
    mma_tiler_mn=(256, 256),
    act_func="dswiglu",
    epilogue_op=None,
)
assert api.check_support()
api.compile()
api.execute(
    a_tensor=a, c_tensor=c, d_row_tensor=d_row, d_col_tensor=d_col,
    sfa_tensor=sfa, padded_offsets=padded_offsets, alpha_tensor=alpha,
    beta_tensor=beta, prob_tensor=prob, dprob_tensor=dprob, dbias_tensor=dbias,
    b_tensor=b, sfb_tensor=sfb,
    sfd_row_tensor=sfd_row, sfd_col_tensor=sfd_col,
    amax_tensor=amax, norm_const_tensor=norm_const,
    current_stream=stream,
)

In the class API, dbias generation is specialized at compile time: if sample_dbias is omitted, dbias_tensor must also be omitted at execute().

Discrete mode:

api = GroupedGemmDgluSm100(
    sample_a=a,
    sample_c=c,
    sample_d_row=d_row,
    sample_d_col=d_col,
    sample_sfa=sfa,
    sample_padded_offsets=padded_offsets,
    sample_alpha=alpha,
    sample_beta=beta,
    sample_prob=prob,
    sample_dprob=dprob,
    # Discrete mode:
    num_experts=num_experts,
    b_shape=(n, k),
    b_dtype=torch.uint8,
    # Configuration
    act_func="dgeglu",
    b_major="k",
    epilogue_op="relu",
)
assert api.check_support()
api.compile()
api.execute(
    a_tensor=a, c_tensor=c, d_row_tensor=d_row, d_col_tensor=d_col,
    sfa_tensor=sfa, padded_offsets=padded_offsets, alpha_tensor=alpha,
    beta_tensor=beta, prob_tensor=prob, dprob_tensor=dprob,
    b_ptrs=b_ptrs, sfb_ptrs=sfb_ptrs,
    current_stream=stream,
)

---

Parameters

Weight Mode

The weight mode is auto-detected from constructor arguments:

• Dense: Provide sample_b and sample_sfb (contiguous weight tensors)

• Discrete: Provide num_experts, b_shape, and b_dtype (per-expert pointer mode)

Providing both or neither raises ValueError.

Input/Output Tensors

• Input tensor A: a_tensor / sample_a

  • Shape: (valid_m, K, 1)

  • Stride: (K, 1, valid_m·K) – must be K-major

  • Dtype (ab_dtype): {float4_e2m1fn_x2, uint8, float8_e4m3fn, float8_e5m2}

• Input tensor B (dense mode): b_tensor / sample_b

  • Shape: (N, K, L) where L = num_groups

  • Stride: K-major or N-major. Must be K-major for FP4.

  • Dtype: Must match A

• Input B pointers (discrete mode): b_ptrs

  • Shape: (num_experts,) – 1-D int64 device tensor of per-expert B data pointers

• Input tensor C: c_tensor / sample_c

  • Shape: (valid_m, 2N, 1) – forward activation with interleaved input/gate

  • Stride: (2N, 1, valid_m·2N) – must be N-major

  • Dtype: {float32, float16, bfloat16, float8_e4m3fn, float8_e5m2}

• Output tensor D_row: d_row_tensor / sample_d_row

  • Shape: (valid_m, 2N, 1)

  • Stride: (2N, 1, valid_m·2N) – must be N-major

  • Dtype (d_dtype): {bfloat16, float32} for FP4; {float8_e4m3fn, float8_e5m2} for FP8

• Output tensor D_col: d_col_tensor / sample_d_col

  • Shape: (valid_m, 2N, 1) – must match D_row dtype and stride

• Input tensor prob: prob_tensor / sample_prob

  • Shape: (valid_m, 1, 1), dtype: float32

• Output tensor dprob: dprob_tensor / sample_dprob

  • Shape: (valid_m, 1, 1), dtype: float32

  • Must be zero-initialized

• Scaling tensors: alpha shape (L,), beta shape (L,), amax shape (L, 2, 1), norm_const shape (1,)

Common Parameters

• acc_dtype: Must be torch.float32

• mma_tiler_mn: Kernel tile size. Default: (256, 256)

  • TILE_M ∈ {128, 256}

  • 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 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"

• 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"

Wrapper-specific Parameters

• d_dtype: Output D tensor data type. Default: torch.bfloat16

• cd_major: Must be "n". Default: "n"

• n (discrete only): B weight N dimension

• b_dtype (discrete only): B weight data type

Wrapper Return Values

Returns a TupleDict (dictionary + tuple unpacking):

• d_row_tensor: Row-quantized dGLU output

• d_col_tensor: Column-quantized dGLU output

• dprob_tensor: Gradient of prob

• amax_tensor: Per-group amax (when d_dtype is bf16/fp16)

• sfd_row_tensor: Row scale factors (when SFD enabled)

• sfd_col_tensor: Column scale factors (when SFD enabled)

---

Support Surface and Constraints

Layouts and Strides

• A must be K-major

• B must be K-major (dense) or K/N-major (discrete). Must be K-major for FP4.

• C, D_row, D_col must be N-major

• All tensors must be 16-byte aligned along the contiguous dimension

Data Types

Format	ab_dtype	sf_dtype	sf_vec_size	d_dtype
MXFP8	float8_e4m3fn or float8_e5m2	{float8_e8m0fnu, float8_e4m3fn}	32	{float8_e4m3fn, float8_e5m2}
NVF4	float4_e2m1fn_x2 or uint8	{float8_e4m3fn, float8_e8m0fnu}	{16, 32}	{bfloat16, float32}

Additional Type Constraints

• A and B must have the same dtype

• Scale factor tensors must have the same dtype

• D_row and D_col must have the same dtype

• dbias must be bfloat16

• sf_dtype=float8_e4m3fn is incompatible with sf_vec_size=32

• FP8 c_dtype with vector_f32=True is not supported

• FP4 ab_dtype only supports d_dtype in {bfloat16, float32}

• For non-dbias paths, FP4 ab_dtype with sf_vec_size=16 and d_dtype=float32 is not supported

• FP8 ab_dtype only supports d_dtype in {float8_e4m3fn, float8_e5m2}

Shapes and Divisibility

• N must be divisible by 32 (32-column blocks for input/gate interleaving)

• Expert count must be <= 1024

• Each group’s M dimension is aligned to m_aligned (256)

• In the class API, dbias is compiled in only when sample_dbias is provided; passing a runtime dbias_tensor without sample_dbias raises ValueError

• use_single_group_runtime_offsets=True is supported only by the block-scaled kernel with exactly one expert. In this mode the kernel derives padded_offsets[0] from runtime A.shape[0] and does not load its value from device memory; the argument must still be an int32 tensor with shape (1,).

Environment

• Requires CUDA with SM100+ compute capability (Blackwell GPUs)

---

Usage Examples

For usage examples, see test cases in test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py (dense mode, unified API) and test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu.py (discrete mode).