Frontend APIFrontend OSS APIsGEMM Fusions

Grouped GEMM + dGLU (SM100)

View as Markdown

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/SM103), with a block-scaled dSwiGLU/dGeGLU path on Rubin (SM107). 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 contractSelected backend
A and B are BF16BF16
matching supported FP4/FP8 A and B plus scale descriptorsblock-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. The block-scaled dGeGLU path supports configurable activation alpha, clamp bounds, and linear offset on both Blackwell and Rubin in dense and discrete weight modes. These values are constructor configuration and remain part of the backward wrapper’s compiled-kernel cache key.

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

Rg=αg2AgBgT,Xg=βgCg.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):

dinput=R prob (gate s),d\mathrm{input} = R\,\mathrm{prob}\,(\mathrm{gate}\,s), dgate=R prob input s(1+gate(1−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

dgate=R prob (clamped_input+linear_offset) s(1+1.702 clamped_gate(1−s)) gate_filter,d\mathrm{gate} = R\,\mathrm{prob}\, (\mathrm{clamped\_input}+\mathrm{linear\_offset})\,s (1 + 1.702\,\mathrm{clamped\_gate}(1-s))\, \mathrm{gate\_filter}, dinput=R prob clamped_gate s input_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

Tg=β1tanh⁡(G/β1)σ(G),Tu=β2tanh⁡(U/β2).T_g=\beta_1\tanh(G/\beta_1)\sigma(G),\qquad T_u=\beta_2\tanh(U/\beta_2).

The fused backward computes

∂Tg∂G=(1−tanh⁡2(G/β1))σ(G)+β1tanh⁡(G/β1)σ(G)(1−σ(G)),\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)), ∂Tu∂U=1−tanh⁡2(U/β2),\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 in dSwiGLU, shape (L,); dGeGLU consumes the saved C directly and ignores beta
    • 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])):

ref[m,n]=αg2∑kdequantize(A[m,k],SFA)⋅dequantize(B[n,k,g],SFB)\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):

  • Deinterleave C into alternating 32-column gate and up blocks. dSwiGLU applies beta_g to C; dGeGLU uses the saved forward values directly.

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"), let gate and up be the original saved C values and ref the GEMM result from Step 1:

gate_clamped = min(gate, glu_clamp_max)
up_clamped = clamp(up, glu_clamp_min, glu_clamp_max)
s = sigmoid(geglu_alpha * gate_clamped)
gate_mask = gate <= glu_clamp_max
up_mask = (glu_clamp_min <= up) & (up <= glu_clamp_max)
dgate = ref * prob * (up_clamped + linear_offset) * s
dgate *= (1 + geglu_alpha * gate_clamped * (1 - s)) * gate_mask
dup = ref * prob * gate_clamped * s * up_mask
dprob += sum(ref * gate_clamped * s * (up_clamped + linear_offset), axis=columns)

The gate has only an upper clamp; the up branch has both bounds. Clamp masks use the original values before clamping, and gradients are retained at equality with either bound. The offset is added after clamping the up branch. The probability gradient does not include a factor of prob, so it can be nonzero when the routing probability is zero. [dgate, dup] is stored in alternating 32-column blocks.

The upstream ref includes alpha_tensor[g] ** 2 on both architectures; this per-group GEMM scale is separate from geglu_alpha. dGeGLU reads C in its stored precision, so a BF16 saved intermediate is the input to this derivative, not the original FP32 forward accumulator. C is not modified.

The defaults are geglu_alpha=1.702, glu_clamp_max=7.0, glu_clamp_min=-7.0, and linear_offset=1.0. For DeepSeek V4 clamped SwiGLU, use act_func="dgeglu", geglu_alpha=1.0, linear_offset=0.0, glu_clamp_max=L, and glu_clamp_min=-L, matching the forward configuration (L=10 for the Flash recipe).

Set these parameters when constructing GroupedGemmDgluSm100 or calling grouped_gemm_dglu_wrapper_sm100. The backward API specializes the activation configuration; changing it selects another compiled-kernel cache entry, while repeating a configuration reuses that entry. It does not expose the forward API’s per-execution activation controls. With BF16 C, the packed derivative represents clamp bounds in BF16; limits such as 7 and 10 are exact, while nonrepresentable limits can differ from the scalar FP32 clamp path.

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

SFD_row[m,n]=norm_const⋅max⁡k∈block∣D[m,k]∣⋅rcp_max\text{SFD\_row}[m, n] = \text{norm\_const} \cdot \max_{k \in \text{block}} |D[m, k]| \cdot \text{rcp\_max} Dquantized[m,n]=D[m,n]⋅norm_constSFD_row[m,n]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]^2
                    v
               ref (valid_m×N×1)
                    |
 C (valid_m×2N×1) --> activation-specific preparation and 32-col split
                    |                    |
                    |                swish, sigmoid
                    |                    |
                    +--> dprob (sum over blocks)
                    |
                    +--> dgate, dup → 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"
  • linear_offset: Offset added to the clamped dGeGLU up branch. Default: 1.0 for dGeGLU
  • geglu_alpha: dGeGLU sigmoid input scale. Default: 1.702
  • glu_clamp_max: dGeGLU upper bound for gate and up. Default: 7.0
  • glu_clamp_min: dGeGLU lower bound for up only. Default: -7.0
  • These activation parameters must match forward and specialize the block-scaled backward cache on Blackwell and Rubin; the BF16 backend retains its fixed alpha/clamp values
  • 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

Formatab_dtypesf_dtypesf_vec_sized_dtype
MXFP8float8_e4m3fn or float8_e5m2{float8_e8m0fnu, float8_e4m3fn}32{float8_e4m3fn, float8_e5m2}
NVF4float4_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/SM103 (Blackwell), or SM107 (Rubin) for the block-scaled backend
  • Rubin MXFP8 uses matching FP8 A/B operands, E8M0 scale factors, and sf_vec_size=32; the dGLU API requires FP8 D outputs with row/column scale factors for these inputs

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). Rubin MXFP8 activation-parameter coverage is in test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py (test_rubin_mxfp8_clamped_dgeglu_*).