> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/cudnn/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/cudnn/_mcp/server.

# Grouped GEMM + WGrad (Unified)

`GroupedGemmWgradSm100` and `grouped_gemm_wgrad_wrapper_sm100` are experimental
SM100+ APIs for grouped MoE weight gradients. The same public surface dispatches
BF16 inputs to the BF16 kernel and preserves the legacy FP4/FP8 block-scaled
backend.

The CuTe DSL dependencies both APIs need ship with the package:

```bash
pip install nvidia-cudnn-frontend
```

## JAX support

Supports **JAX arrays** on the BF16 backend: A k-major and B n-major C-contiguous arrays, dense `(experts, m, n)` C-contiguous output or discrete output pointers (packed uint8 / int64 with jax x64 mode). The block-scaled backend's layouts are not expressible as JAX arrays and raise a clear error. The wrapper is eager, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs.

For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `grouped_gemm_wgrad_jax_sm100` (built on `cudnn.jax.call`; BF16, discrete output pointers): the per-expert weight-gradient buffers behind `wgrad_ptrs` are caller-owned external memory the kernel writes through, so the entry returns a completion **token** — `jax.block_until_ready(token)` before reading them (and zero them yourself between runs unless accumulating). Under tracing the per-group offsets *values* cannot be host-validated; the external buffers must stay alive and unmoved across every execution of the traced computation.

## Operation

For expert `e`, let `begin = 0` for the first expert and
`begin = offsets_tensor[e - 1]` otherwise, and let
`end = offsets_tensor[e]`. The API computes

```text
Wgrad[e] = A[:, begin:end] @ B[begin:end, :]
```

When `accumulate_on_output=True`, that result is accumulated into the existing
output. The caller must therefore initialize every output allocation. When it
is false, the kernel overwrites the output; an empty expert produces zero.

## BF16 contract

The BF16 backend accepts:

| Argument | Shape | Supported stride/major | Dtype |
| --- | --- | --- | --- |
| `a_tensor` | `(hidden, tokens_sum)` | `(tokens_sum, 1)` K-major or `(1, hidden)` M-major | `torch.bfloat16` |
| `b_tensor` | `(tokens_sum, intermediate)` | `(1, tokens_sum)` K-major or `(intermediate, 1)` N-major | `torch.bfloat16` |
| `offsets_tensor` | `(num_experts,)` | contiguous `(1,)` | `torch.int32` |
| dense `wgrad_tensor` | `(num_experts, hidden, intermediate)` | `(hidden * intermediate, intermediate, 1)` | BF16, FP16, or FP32 |
| one discrete output | `(hidden, intermediate)` | `(intermediate, 1)` | BF16, FP16, or FP32 |
| `wgrad_ptrs` | `(num_experts,)` | contiguous `(1,)` | `torch.int64` |

`offsets_tensor` is a non-decreasing cumulative sum. Every expert token count
(`offsets[e] - offsets[e - 1]`) must be a multiple of 256, and the final offset
must equal `tokens_sum`. Inputs, metadata, and outputs must reside on the same
CUDA device and satisfy the API's alignment checks.

BF16 uses FP32 accumulation and requires `sf_vec_size=16`. Pass `None` for
`sfa_tensor`, `sfb_tensor`, `global_scale_a`, and `global_scale_b`. BF16 rejects
every non-`None` scale or global-scale control with `ValueError`; it never falls
through to another backend. Only a supported FP4/FP8 operand pair selects the
legacy block-scaled backend, which continues to support its existing scale
tensors and global scales.

`input_order` describes how the token dimension is stored:

- `"tensor2d"` (default) uses one global 2-D tensor and its declared strides.
- `"tensor_ragged"` uses per-expert K-contiguous blocks concatenated in memory.
  In this mode only each input's unit-stride axis is meaningful; non-unit host
  strides are ignored when per-expert TMA descriptors are built.

### Output modes

With `output_mode="dense"`, provide or let the wrapper allocate the contiguous
stacked `wgrad_tensor`. `wgrad_ptrs` is forbidden.

With `output_mode="discrete"`, either:

- omit both output arguments and let the wrapper allocate a stacked tensor and
  construct an internal pointer array; or
- provide a CUDA `torch.int64` `wgrad_ptrs` array containing one non-null,
  16-byte-aligned output address per expert.

For explicit pointer-only output, `result["wgrad_tensor"]` is `None`. The caller
owns all pointed-to output allocations and must keep both those allocations and
the pointer tensor alive until work on `current_stream` completes. The API
records the pointer tensor on the launch stream, but it cannot manage the
lifetime of allocations represented only by integer addresses.

The wrapper always returns `TupleDict(wgrad_tensor=...)`; it contains exactly
one item and supports either keyed access or tuple unpacking.

## Block-scaled contract

The legacy block-scaled backend is selected only by a supported matching FP4/FP8
operand pair. It preserves the pre-existing scale-factor contract: provide
`sfa_tensor` and `sfb_tensor`, and provide `global_scale_a` and
`global_scale_b` where the selected low-precision format requires them. BF16
does not reinterpret these controls; it rejects them instead.

Torch block-scaled callers in dense or discrete output mode that retain
operations for CUDA Graph replay may provide a caller-owned
`descriptor_workspace`; pass this extension by keyword. Allocate its size with
`get_grouped_gemm_wgrad_workspace_size_sm100`, keep it alive for as long as the
captured call site may replay, and do not share it between call sites that may
overlap. This lets multiple same-signature calls share one compiled kernel
without sharing mutable runtime TMA descriptors. Callers that omit this
argument retain the compatibility behavior that isolates cached API instances
by explicit dense output address; discrete callers retain the compiled
operation's internal workspace.

```python
workspace = torch.empty(
    cudnn.get_grouped_gemm_wgrad_workspace_size_sm100(num_experts),
    dtype=torch.uint8,
    device=a_tensor.device,
)
result = cudnn.grouped_gemm_wgrad_wrapper_sm100(
    a_tensor=a_tensor,
    b_tensor=b_tensor,
    sfa_tensor=sfa_tensor,
    sfb_tensor=sfb_tensor,
    offsets_tensor=offsets_tensor,
    wgrad_tensor=wgrad_tensor,
    descriptor_workspace=workspace,
    output_mode="dense",
)
```

## API usage

### BF16

#### Wrapper

Dense BF16 output:

```python
import cudnn
import torch

result = cudnn.grouped_gemm_wgrad_wrapper_sm100(
    a_tensor=a_tensor,
    b_tensor=b_tensor,
    sfa_tensor=None,
    sfb_tensor=None,
    offsets_tensor=offsets_tensor,
    output_mode="dense",
    wgrad_dtype=torch.bfloat16,
    input_order="tensor2d",
)
wgrad_tensor = result["wgrad_tensor"]
```

Discrete BF16 outputs owned by the caller:

```python
expert_outputs = [
    torch.empty(
        (hidden, intermediate), dtype=torch.bfloat16, device="cuda"
    )
    for _ in range(offsets_tensor.numel())
]
wgrad_ptrs = torch.tensor(
    [output.data_ptr() for output in expert_outputs],
    dtype=torch.int64,
    device="cuda",
)

result = cudnn.grouped_gemm_wgrad_wrapper_sm100(
    a_tensor=a_tensor,
    b_tensor=b_tensor,
    sfa_tensor=None,
    sfb_tensor=None,
    offsets_tensor=offsets_tensor,
    output_mode="discrete",
    wgrad_ptrs=wgrad_ptrs,
    wgrad_dtype=torch.bfloat16,
    input_order="tensor_ragged",
)
assert result["wgrad_tensor"] is None
```

#### Reusable class lifecycle

The class API requires output descriptors at construction and output storage at
execution. This dense BF16 example compiles once and accepts later calls with a
different `tokens_sum` when static dimensions, dtypes, majors, and configuration
remain compatible:

```python
op = cudnn.GroupedGemmWgradSm100(
    sample_a=a_tensor,
    sample_b=b_tensor,
    sample_sfa=None,
    sample_sfb=None,
    sample_offsets=offsets_tensor,
    sample_wgrad=wgrad_tensor,
    acc_dtype=torch.float32,
    input_order="tensor2d",
)
op.check_support()
op.compile()
op.execute(
    a_tensor=a_tensor,
    b_tensor=b_tensor,
    sfa_tensor=None,
    sfb_tensor=None,
    offsets_tensor=offsets_tensor,
    wgrad_tensor=wgrad_tensor,
)
```

For a discrete class instance, replace `sample_wgrad` with
`sample_wgrad_expert=expert_outputs[0]`, `num_experts`, `wgrad_shape`, and
`wgrad_dtype`, then pass `wgrad_ptrs` to `execute`.

### Block-scaled

#### Wrapper

```python
result = cudnn.grouped_gemm_wgrad_wrapper_sm100(
    a_tensor=a_tensor,
    b_tensor=b_tensor,
    sfa_tensor=sfa_tensor,
    sfb_tensor=sfb_tensor,
    offsets_tensor=offsets_tensor,
    output_mode="dense",
    wgrad_dtype=torch.bfloat16,
    input_order="tensor_ragged",
)
```

#### Reusable class lifecycle

```python
op = cudnn.GroupedGemmWgradSm100(
    sample_a=a_tensor,
    sample_b=b_tensor,
    sample_sfa=sfa_tensor,
    sample_sfb=sfb_tensor,
    sample_offsets=offsets_tensor,
    sample_wgrad=sample_wgrad_tensor,
    acc_dtype=torch.float32,
)
assert op.check_support()
op.compile()
op.execute(
    a_tensor=a_tensor,
    b_tensor=b_tensor,
    sfa_tensor=sfa_tensor,
    sfb_tensor=sfb_tensor,
    offsets_tensor=offsets_tensor,
    wgrad_tensor=wgrad_tensor,
)
```

## Scheduling, cache, and errors

The BF16 kernel uses dynamic persistent scheduling. The token dimension is
compiled dynamically, and the wrapper cache abstracts the token-sized axes of A
and B while retaining static dimensions, layouts, dtypes, output descriptors,
tiling, cluster shape, input order, and accumulation mode in its key. A changed
static contract creates a different cached operator or fails validation.

The APIs reject unsupported dtypes or layouts, malformed/unaligned offsets or
pointers, mixed devices, forbidden BF16 scale controls, unsupported tiling, use
before `compile()`, unavailable CUDA, and devices below SM100. Support and
validation errors are reported as `ValueError` or `RuntimeError`; callers should
not rely on this experimental API remaining source-compatible across releases.