feat: add BF16 grouped GEMM MoE kernels

.pre-commit-config.yaml

@@ -15,7 +15,7 @@ repos:
         minimum_pre_commit_version: 2.9.2
         require_serial: true
         types_or: [python, pyi]
-        args: ["--line-length", "160"]
+        args: ["--line-length", "160", "--fast"]
     -   id: black-jupyter
         name: black-jupyter
         description:

README.md

@@ -23,9 +23,10 @@ We are now shipping **OSS kernels**, allowing you to inspect, modify, and contri
 *   **[GEMM + SwiGLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/gemm_swiglu):** High-performance implementation of the SwiGLU activation fused with GEMM.
 *   **[GEMM + sReLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/gemm_srelu):** High-performance implementation of squared-ReLU fused with GEMM.
 *   **[GEMM + dsReLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/gemm_dsrelu):** High-performance implementation of dsquared-ReLU fused with GEMM.
-*   **[Grouped GEMM + GLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_glu):** Unified grouped GEMM GLU API supporting dense and discrete MoE weight layouts.
+*   **[Grouped GEMM (BF16)](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_unfused):** Unfused BF16 grouped GEMM with dense and discrete MoE weight layouts.
+*   **[Grouped GEMM + GLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_glu):** Unified BF16 and legacy block-scaled grouped GEMM GLU API supporting dense and discrete MoE weight layouts.
 *   **[Grouped GEMM + GLU + Hadamard](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard):** Dense grouped GEMM GLU forward fusion with a fused Hadamard transform and per-expert AMAX reduction.
-*   **[Grouped GEMM + dGLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_dglu):** Unified grouped GEMM dGLU backward API supporting dense and discrete MoE weight layouts.
+*   **[Grouped GEMM + dGLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_dglu):** Unified BF16 and legacy block-scaled grouped GEMM dGLU backward API supporting dense and discrete MoE weight layouts.
 *   **[Grouped GEMM + SwiGLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_swiglu):** SwiGLU activation fused with Grouped GEMM.
 *   **[Grouped GEMM + dSwiglu](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_dswiglu):** dSwiglu activation fused with Grouped GEMM.
 *   **[Grouped GEMM + sReLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_srelu):** Contiguous grouped squared-ReLU GEMM for MoE workloads.
@@ -34,7 +35,7 @@ We are now shipping **OSS kernels**, allowing you to inspect, modify, and contri
 *   **[Discrete Grouped GEMM + dSwiGLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu):** Per-expert-pointer dSwiGLU backward grouped GEMM for MoE workloads without weight packing.
 *   **[Grouped GEMM + Quant](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_quant):** Legacy dense-only grouped GEMM quant API for MoE FC2/dFC1 workloads.
 *   **[Grouped GEMM + Quant (Unified)](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_quant):** Unified grouped GEMM quant API with per-row gating for MoE FC2/dFC1 workloads.
-*   **[Grouped GEMM + Wgrad](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_wgrad):** Unified grouped GEMM weight-gradient API supporting dense and discrete output layouts for MoE workloads.
+*   **[Grouped GEMM + Wgrad](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_wgrad):** Unified BF16 and legacy block-scaled grouped GEMM weight-gradient API supporting dense and discrete output layouts for MoE workloads.
 *   **[BSA](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/block_sparse_attention/):** Block-sparse attention forward and backward CuTe DSL kernels for block-level routing metadata.
 *   **[NSA](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/native_sparse_attention/):** Native Sparse attention as described in the Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention.
 *   **[SDPA Backward: SM100, D=256](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/sdpa):** SDPA Backward pass for D=256 on SM100.

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

@@ -0,0 +1,153 @@
+# Grouped GEMM (SM100 BF16)
+
+**This is an experimental API and subject to change.** It requires an NVIDIA
+SM100-or-newer GPU and the optional CuTe DSL dependencies:
+
+```bash
+pip install nvidia-cudnn-frontend[cutedsl]
+```
+
+`GroupedGemmSm100` and `grouped_gemm_wrapper_sm100` implement the neutral,
+unfused BF16 MoE grouped GEMM. They support dense stacked expert weights or a
+device pointer per expert, optional bias, optional materialization of the
+intermediate `C`, and static or dynamic tile scheduling.
+
+## Operation
+
+Let `padded_offsets[g]` be the exclusive end of expert `g`'s contiguous row
+range and let `begin_g` be zero for the first expert or
+`padded_offsets[g - 1]` otherwise. For rows in `[begin_g, padded_offsets[g])`:
+
+```text
+G_g = alpha[g] * A_g @ B_g.T
+C_g = G_g + prob_g * bias[:, g]     # when bias is present
+D_g = C_g
+```
+
+Without bias, `C_g = G_g` and `D_g = prob_g * G_g`. Accumulation is FP32;
+`C` and `D` may independently use BF16, FP16, or FP32.
+
+## Tensors and layouts
+
+For total padded rows `M`, inner dimension `K`, output dimension `N`, and `L`
+experts:
+
+| Tensor | Shape | Required stride / dtype |
+| --- | --- | --- |
+| `A` | `(M, K, 1)` | `(K, 1, M*K)`, BF16 |
+| dense `B` | `(N, K, L)` | `(K, 1, N*K)`, BF16 |
+| discrete `b_ptrs` | `(L,)` | contiguous CUDA int64 pointers to `(N, K)` BF16 matrices |
+| `padded_offsets` | `(L,)` | `(1,)`, CUDA int32 cumulative ends |
+| `alpha` | `(L,)` | `(1,)`, CUDA FP32 |
+| `prob` | `(M, 1, 1)` | `(1, 1, 1)`, CUDA FP32 |
+| optional `bias` | `(N, L)` | `(1, N)`, BF16/FP16/FP32 |
+| `C`, `D` | `(M, N, 1)` | `(N, 1, M*N)`, BF16/FP16/FP32 |
+
+`M` and every cumulative offset must be 256-aligned. Dense weights are
+K-major. Discrete mode uses `b_major="k"`, `n=N`, and
+`b_dtype=torch.bfloat16`.
+
+The pointer-array tensor must be contiguous, non-null, eight-byte aligned, and
+on the same device as `A`; each target pointer must satisfy the kernel's
+alignment contract. The API records the pointer-array tensor on the launch
+stream. The caller must keep every pointed-to expert allocation alive and must
+not modify or free it until that stream completes.
+
+## Wrapper API
+
+Dense mode:
+
+```python
+import cudnn
+import torch
+
+result = cudnn.grouped_gemm_wrapper_sm100(
+    a_tensor=a,
+    padded_offsets=padded_offsets,
+    alpha_tensor=alpha,
+    b_tensor=b,
+    bias_tensor=bias,
+    prob_tensor=prob,
+    c_dtype=torch.float32,
+    d_dtype=torch.bfloat16,
+    generate_c=True,
+    use_dynamic_sched=True,
+)
+d, c = result
+assert d is result["d_tensor"]
+assert c is result["c_tensor"]
+```
+
+Discrete mode changes only the weight arguments:
+
+```python
+result = cudnn.grouped_gemm_wrapper_sm100(
+    a_tensor=a,
+    padded_offsets=padded_offsets,
+    alpha_tensor=alpha,
+    b_ptrs=b_ptrs,
+    n=N,
+    b_dtype=torch.bfloat16,
+    b_major="k",
+    bias_tensor=bias,
+    prob_tensor=prob,
+)
+```
+
+The `TupleDict` order is exactly `d_tensor`, then `c_tensor`. `c_tensor` is
+`None` unless `generate_c=True`; the kernel still uses an internal C buffer
+when it is needed for execution.
+
+## Class API
+
+The class API takes representative sample tensors, then follows
+`check_support()` -> `compile()` -> `execute()`:
+
+```python
+op = cudnn.GroupedGemmSm100(
+    sample_a=a,
+    sample_c=c,
+    sample_d=d,
+    sample_padded_offsets=padded_offsets,
+    sample_alpha=alpha,
+    sample_b=b,
+    sample_bias=bias,
+    sample_prob=prob,
+    acc_dtype=torch.float32,
+    generate_c=True,
+    use_dynamic_sched=False,
+)
+assert op.check_support()
+op.compile()
+op.execute(
+    a_tensor=a,
+    c_tensor=c,
+    d_tensor=d,
+    padded_offsets=padded_offsets,
+    alpha_tensor=alpha,
+    b_tensor=b,
+    bias_tensor=bias,
+    prob_tensor=prob,
+)
+```
+
+For a discrete class instance, replace `sample_b` with
+`num_experts=L`, `b_shape=(N, K)`, `b_dtype=torch.bfloat16`, and pass
+`b_ptrs` to `execute()`.
+
+## Scheduling, caching, and errors
+
+- `use_dynamic_sched=False` uses the static scheduler.
+- `use_dynamic_sched=True` compiles a dynamic-M callable and reuses it for
+  compatible M values; discrete mode also allocates the per-expert tensor-map
+  workspace. Wrapper cache keys retain dtype, layout, expert count, optional
+  features, scheduler choice, output policy, tile/cluster shape, and overlap
+  margin.
+- Dense and discrete weight arguments are mutually exclusive. Invalid shapes,
+  strides, dtypes, devices, alignment, offsets, pointer entries, output
+  descriptors, tiles/clusters, or a target below SM100 raise `ValueError` or
+  `RuntimeError` before launch.
+- Fused GLU, dGLU, and WGrad APIs select BF16 from BF16 operands while keeping
+  their existing FP4/FP8 block-scaled backends. For BF16 on those fused APIs,
+  scale-factor controls are `None`; see their operation pages for the exact
+  dispatch and return contracts.

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

@@ -4,7 +4,10 @@
 
 ## Overview
 
-**Unified Grouped GEMM + dGLU fusion**: A block-scaled grouped GEMM fused with a dGLU backward epilogue (dSwiGLU or dGeGLU) on NVIDIA Blackwell GPUs (SM100+), designed for MoE (Mixture of Experts) workloads. Implemented with CUTLASS/CUTE.
+**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
+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
@@ -16,7 +19,98 @@ And both backward activation functions:
 
 Groups are contiguous in the M dimension and described by `padded_offsets` (cumulative aligned end offsets).
 
-This kernel performs:
+### 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.
+
+## 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:
+
+```text
+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
@@ -110,9 +204,74 @@ $$
 
 ---
 
-## API Usage
+## API usage
+
+### BF16
+
+#### High-level wrapper
+
+```python
+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
+
+```python
+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
+#### High-level wrapper
 
 **Dense mode:**
 
@@ -185,7 +344,7 @@ outputs = grouped_gemm_dglu_wrapper_sm100(
 )
 ```
 
-### Class API
+#### Class API
 
 **Dense mode:**