NVIDIA cuDNN

https://docs.nvidia.com/deeplearning/cudnn/latest
Admin
NVIDIA cuDNN is a GPU-accelerated library of primitives for deep neural networks, providing highly...
Tokens:112,340
Snippets:1,045
Trust Score:10
Update:2 months ago
Show doc for...
Tokens:
Raw
### Install CuTe DSL dependencies

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/bsa.html

Install the necessary optional dependencies to use the CuTe DSL kernels.

```bash
pip install nvidia-cudnn-frontend[cutedsl]
```

--------------------------------

### Implement autograd for custom operations

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/utilities/adding_torch_custom_ops.html

Define context setup and backward functions, then register them with the PyTorch library.

```python
def _setup_context(ctx, inputs, output):
    x, w, eps, training, bias = inputs
    y, aux = output
    ctx.save_for_backward(x, w, y, aux)
    ctx.eps = eps
    # save non-tensor args as ctx attributes

def _backward(ctx, dY, d_aux):
    x, w, y, aux = ctx.saved_tensors
    dX, dW = torch.ops.cudnn.my_op_bwd(dY, x, w, y, aux, ctx.eps)
    return dX, dW, None, None, None  # None for non-differentiable args

torch.library.register_autograd("cudnn::my_op", _backward, setup_context=_setup_context)
```

--------------------------------

### Ragged Layout Packing Example

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/operations/Attention.html

Shows how sequences are packed together without padding in a ragged layout.

```text
Q = aabbb  (no padding, T=5 total tokens)
```

--------------------------------

### Install Python Dependencies

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/samples.html

Command to install required Python packages for running cuDNN samples.

```bash
pip install -r requirements.txt
```

--------------------------------

### Partially Packed Layout Example

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/operations/Attention.html

Illustrates tokens within each batch being contiguous without being globally packed.

```text
Q = aa00bbb0  (batch 0 at offset 0, batch 1 at offset 4)
```

--------------------------------

### Padding Mask Sequence Example

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/operations/Attention.html

Demonstrates the structure of query tensors when using padding masks for sequences of different lengths.

```text
Q[b=0] = aa000000  (6 padding tokens)
Q[b=1] = bbb00000  (5 padding tokens)
```

--------------------------------

### Build RoPE and SDPA Graph

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/operations/RoPE.html

Example of constructing a graph with RoPE operations for Q and K tensors followed by an SDPA operation.

```python
q_rot = graph.rope(input=q, freqs=freqs, output_scale=α_q, rope_dim=...)
k_rot = graph.rope(input=k, freqs=freqs, output_scale=α_k, rope_dim=...)
o, stats = graph.sdpa(q=q_rot, k=k_rot, v=v, attn_scale=...)
```

--------------------------------

### Forward MoE Grouped Matmul Example (Python)

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/operations/MoeGroupedMatmul.html

Demonstrates setting up a forward MoE grouped matmul operation using the cuDNN pygraph API.

```python
import torch
import cudnn

num_experts = 8
token_num   = 1024   # S
hidden_size = 512    # K
weight_size = 256    # N

graph = cudnn.pygraph(
    intermediate_data_type=cudnn.data_type.FLOAT,
    compute_data_type=cudnn.data_type.FLOAT,
)

# Token: [1, S, K], bf16, row-major
token_t = graph.tensor(
    name="token",
    dim=[1, token_num, hidden_size],
    stride=[token_num * hidden_size, hidden_size, 1],
    data_type=cudnn.data_type.BFLOAT16,
)

# Weight: [E, K, N], bf16, column-major inner dims (stride[1] == 1)
weight_t = graph.tensor(
    name="weight",
    dim=[num_experts, hidden_size, weight_size],
    stride=[hidden_size * weight_size, 1, hidden_size],
    data_type=cudnn.data_type.BFLOAT16,
)

# FirstTokenOffset: [E, 1, 1], INT32
fto_t = graph.tensor(
    name="first_token_offset",
    dim=[num_experts, 1, 1],
    stride=[1, 1, 1],
    data_type=cudnn.data_type.INT32,
)

output_t = graph.moe_grouped_matmul(
    token_t, weight_t, fto_t,
    mode=cudnn.moe_grouped_matmul_mode.NONE,
    compute_data_type=cudnn.data_type.FLOAT,
    name="moe_fwd",
)
output_t.set_output(True).set_data_type(cudnn.data_type.BFLOAT16)

graph.validate()
graph.build_operation_graph()
graph.create_execution_plans([cudnn.heur_mode.A])
graph.check_support()
graph.build_plans()
```

--------------------------------

### Constructing the RoPE backward graph

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/operations/RoPE.html

Example of building the backward graph for RoPE operations using the provided graph API.

```python
dq_rot, dk_rot, dv = graph.sdpa_backward(q=q_rot, k=k_rot, v=v, o=o, dO=dO, stats=stats, attn_scale=...)
dq = graph.rope_backward(dY=dq_rot, freqs=freqs, output_scale=α_q, rope_dim=...)
dk = graph.rope_backward(dY=dk_rot, freqs=freqs, output_scale=α_k, rope_dim=...)
```

--------------------------------

### Backward MoE Grouped Matmul Example (Python)

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/operations/MoeGroupedMatmul.html

Demonstrates setting up a backward MoE grouped matmul operation using the cuDNN pygraph API.

```python
import torch
import cudnn

num_experts = 8
token_num   = 1024
hidden_size = 512
weight_size = 256

graph = cudnn.pygraph(
    intermediate_data_type=cudnn.data_type.FLOAT,
    compute_data_type=cudnn.data_type.FLOAT,
)

# dOutput: [1, S, N], bf16, row-major
doutput_t = graph.tensor(
    name="doutput",
    dim=[1, token_num, weight_size],
    stride=[token_num * weight_size, weight_size, 1],
    data_type=cudnn.data_type.BFLOAT16,
)

# Token: [1, S, K], bf16, row-major
token_t = graph.tensor(
    name="token",
    dim=[1, token_num, hidden_size],
    stride=[token_num * hidden_size, hidden_size, 1],
    data_type=cudnn.data_type.BFLOAT16,
)

# FirstTokenOffset: [E, 1, 1], INT32
fto_t = graph.tensor(
    name="first_token_offset",
    dim=[num_experts, 1, 1],
    stride=[1, 1, 1],
    data_type=cudnn.data_type.INT32,
)

dweight_t = graph.moe_grouped_matmul_bwd(
    doutput_t, token_t, fto_t,
    compute_data_type=cudnn.data_type.FLOAT,
    name="moe_bwd",
)
# dweight: [E, K, N]
dweight_t.set_output(True).set_data_type(cudnn.data_type.BFLOAT16)

graph.validate()
graph.build_operation_graph()
graph.create_execution_plans([cudnn.heur_mode.A])
graph.check_support()
graph.build_plans()
```

--------------------------------

### Compile cuDNN Application

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/quickstart.html

Example command to compile a cuDNN application using nvcc.

```bash
nvcc -std=c++17 -I path/to/cudnn_frontend/include -I path/to/cudnn/include -L path/to/cudnn/lib -lcudnn -lnvrtc -o simple_matmul simple_matmul.cpp
```

--------------------------------

### SDPA Example Usage

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/operations/Attention.html

Demonstrates forward pass and autograd-based backward pass using the cuDNN SDPA operator.

```python
import torch
from cudnn.experimental.ops import scaled_dot_product_attention

B, H, S, D = 2, 8, 1024, 128

q = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda", requires_grad=True)
k = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda", requires_grad=True)
v = torch.randn(B, H, S, D, dtype=torch.float16, device="cuda", requires_grad=True)

# Forward
output = scaled_dot_product_attention(q, k, v, is_causal=True)

# Backward (autograd handles this automatically)
loss = output.sum()
loss.backward()
# q.grad, k.grad, v.grad are now populated
```

--------------------------------

### Unsupported Interleaved Token Layout

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/operations/Attention.html

Example of a non-contiguous token layout that is not supported by the ragged layout implementation.

```text
Q = a0abbb00bb000000  (tokens interleaved - NOT SUPPORTED)
```

--------------------------------

### Full NSA Pipeline Implementation

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/nsa.html

Demonstrates the complete NSA workflow including compression, top-k reduction, selection attention, and sliding window attention.

```python
import torch
from cudnn import NSA
import math

# Configuration
batch_size = 2
seq_len = 2048
h_q, h_kv = 32, 8  # GQA with 4:1 ratio
head_dim = 128
k_value = 16
block_size = 64
compress_stride = 32
window_size = 512

device = "cuda"
dtype = torch.bfloat16

# Allocate input tensors (T,H,D format)
seq_lengths = torch.tensor([seq_len] * batch_size, dtype=torch.int32, device=device)
cum_seqlen = torch.tensor([0, seq_len, 2 * seq_len], dtype=torch.int32, device=device)
T = cum_seqlen[-1].item()

q = torch.randn(T, h_q, head_dim, dtype=dtype, device=device)
k = torch.randn(T, h_kv, head_dim, dtype=dtype, device=device)
v = torch.randn(T, h_kv, head_dim, dtype=dtype, device=device)

# Step 1: Compression Attention
o_cmp, lse_cmp = NSA.compression_attention_wrapper(
    q_tensor=q,
    k_tensor=k,
    v_tensor=v,
    cum_seqlen_q_tensor=cum_seqlen,
    cum_seqlen_k_tensor=cum_seqlen,
    enable_lse=True,
    o_dtype=torch.bfloat16,
)

# Step 2: Top-K Reduction to find important blocks (uses LSE from compression)
topk_scores, topk_indices = NSA.topk_reduction_wrapper(
    q_tensor=q,
    k_tensor=k,
    lse_tensor=lse_cmp,
    cum_seqlen_q_tensor=cum_seqlen,
    cum_seqlen_k_tensor=cum_seqlen,
    k_value=k_value,
    selection_block_size=block_size,
    compress_stride=compress_stride,
    is_causal=True,
)

# Step 3: Selection Attention for important blocks
block_counts = torch.full((T, h_kv), k_value, dtype=torch.int32, device=device)

o_sel, l_sel, m_sel = NSA.selection_attention_wrapper(
    q_tensor=q,
    k_tensor=k,
    v_tensor=v,
    block_indices_tensor=topk_indices,
    block_counts_tensor=block_counts,
    cum_seqlen_q_tensor=cum_seqlen,
    cum_seqlen_k_tensor=cum_seqlen,
    block_size=block_size,
)

# Step 4: Sliding Window Attention for local context
o_swa, stats_swa = NSA.sliding_window_attention_wrapper(
    q_tensor=q,
    k_tensor=k,
    v_tensor=v,
    seq_len_q_tensor=seq_lengths,
    seq_len_kv_tensor=seq_lengths,
    left_bound=window_size,
    right_bound=0,
    is_infer=False,
    attn_scale=1.0 / math.sqrt(head_dim),
)

# Step 5: Combine outputs (application-specific weighted combination)
# This is a simplified example; actual combination uses learned gating
final_output = o_cmp + o_sel + o_swa  # Placeholder combination
```

--------------------------------

### Build C++ Samples

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/samples.html

Commands to configure and compile the C++ samples using CMake.

```bash
mkdir build
cd build
cmake -DCUDNN_PATH=/path/to/cudnn -DCUDAToolkit_ROOT=/path/to/cuda  ../
cmake --build . -j16
bin/samples
```

--------------------------------

### Set Library Path and Run Executable

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/quickstart.html

Configures the library search path and executes the compiled binary.

```bash
# optional to set LD_LIBRARY_PATH
export LD_LIBRARY_PATH=/path/to/cudnn/lib
# run the executable
./simple_matmul
```

--------------------------------

### Python Slice Operation Example

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/operations/Slice.html

Demonstrates how to apply a slice operation to a tensor using the Python API.

```python
input_tensor = graph.tensor(dim=[4, 8, 16], stride=[128, 16, 1], data_type=cudnn.data_type.float)
sliced_tensor = graph.slice(
    input_tensor,
    slices=[slice(1, 3), slice(2, 6, 2), slice(0, 16)],
    name="my_slice",
    compute_data_type=cudnn.data_type.float,
)
```

--------------------------------

### Initialize Device Properties

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/utilities/deviceless-ahead-of-time-compilation.html

Methods for populating the device property descriptor from a handle, device ID, or serialized buffer.

```cpp
set_handle(cudnnHandle_t handle);  // initialize from a cuDNN handle
set_device_id(int32_t device_id);  // initialize from a specific device
deserialize(const std::vector<uint8_t>& serialized_buf);  // deserialize from json
```

--------------------------------

### Create Custom Execution Plan

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/utilities/custom-execution-plan.html

Initializes an execution plan using a specified engine ID and a map of knob configurations.

```cpp
error_t
create_execution_plan(int64_t const engine_id, std::unordered_map<KnobType_t, int64_t> const &knobs);
```

--------------------------------

### Create and Set Custom Kernel Cache

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/developer/graph-api.html

Demonstrates the steps to create a custom kernel cache, enable dynamic shapes, and associate the cache with a graph. This is useful for managing cache lifetime and sharing caches across specific plans.

```cpp
auto kernel_cache = std::make_shared<cudnn_frontend::KernelCache>();
graph.set_dynamic_shape_enabled(true);
graph.set_kernel_cache(kernel_cache);
```

--------------------------------

### Get Execution Plan Count

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/developer/overview.html

Returns the total number of execution plans available, where index 0 represents the highest priority.

```cpp
cudnn_frontend::int64_t
cudnn_frontend::Graph::get_execution_plan_count() const;
```

--------------------------------

### Create a Device Property Descriptor

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/utilities/deviceless-ahead-of-time-compilation.html

Initializes a shared pointer for the device properties descriptor.

```cpp
auto device_prop = std::make_shared<cudnn_frontend::DeviceProperties>();
```

--------------------------------

### Initialize and Execute SDPA Forward Class API

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/attention/sdpa_fwd_d256.html

Uses the class-based API to instantiate, compile, and execute the SDPA forward pass. Requires sample tensors for initialization to determine memory requirements.

```python
import torch
from cudnn import SdpafwdSm100D256

sdpa_fwd = SdpafwdSm100D256(
    sample_q=q,
    sample_k=k,
    sample_v=v,
    sample_o=o,
    sample_lse=lse,
    sample_cum_seqlen_q=cum_seqlen_q,  # None for B,H,S,D layout
    sample_cum_seqlen_k=cum_seqlen_k,  # None for B,H,S,D layout
    max_s_q=1024,
    max_s_k=1024,
    qk_acc_dtype=torch.float32,
    pv_acc_dtype=torch.float32,
    mma_tiler_mn=(128, 128),
    is_causal=False,
    window_size=(-1, -1),
    scale_softmax=None,
    scale_output=1.0,
)

assert sdpa_fwd.check_support()
sdpa_fwd.compile()
sdpa_fwd.execute(
    q_tensor=q,
    k_tensor=k,
    v_tensor=v,
    o_tensor=o,
    lse_tensor=lse,
    cum_seqlen_q_tensor=cum_seqlen_q,
    cum_seqlen_k_tensor=cum_seqlen_k,
    scale_softmax=None,
    scale_output=1.0,
    current_stream=None,
)
```

--------------------------------

### High-level Wrapper Usage

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.html

Demonstrates the functional wrapper for Grouped GEMM DGLU in dense and discrete modes.

```python
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
```

```python
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,
)
```

--------------------------------

### Execute Indexer Backward Wrapper

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/dsa.html

Performs the three-stage sparse top-K pipeline to produce training gradients. Raises RuntimeError if the CuTe-DSL path fails.

```python
result = DSA.indexer_backward_wrapper(
    index_q, weights, index_k,
    attn_score, index_score, topk_indices,
    sm_scale=1.0, loss_coeff=1.0, grad_loss=1.0, block_I=128,
)
d_index_q, d_weights, d_index_k = (
    result["d_index_q"], result["d_weights"], result["d_index_k"],
)
```

--------------------------------

### Execute Grouped GEMM Wgrad via Wrapper

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.html

Uses the convenience wrapper to perform the operation. The wrapper automatically handles zero-initialization of the output tensor.

```python
import cudnn
import torch

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",
)

wgrad_tensor = result["wgrad_tensor"]
```

--------------------------------

### Registering a Custom Op with torch.Library

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/utilities/adding_torch_custom_ops.html

Use the direct torch.Library API to define schemas, implement CUDA kernels, and register fake/autograd functions for optimal performance.

```python
import torch

_lib = torch.library.Library("cudnn", "DEF")

# Define schema — must list all tensor and scalar args with types
_lib.define(
    "my_op(Tensor x, Tensor w, float eps, bool training=False, "
    "Tensor? bias=None) -> (Tensor, Tensor)"
)

# Register CUDA implementation
def _my_op_impl(x, w, eps, training=False, bias=None):
    # ... build/cache graph, execute ...
    return output, aux

_lib.impl("my_op", _my_op_impl, "CUDA")

# Register fake (for torch.compile shape inference)
@torch.library.register_fake("cudnn::my_op")
def _my_op_fake(x, w, eps, training=False, bias=None):
    return torch.empty_like(x), torch.empty(...)

# Register autograd (if backward is needed)
torch.library.register_autograd("cudnn::my_op", _my_op_backward, setup_context=_my_op_setup_ctx)
```

--------------------------------

### discrete_grouped_gemm_swiglu_wrapper_sm100

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.html

A high-level functional wrapper to execute a discrete grouped GEMM with SwiGLU activation.

```APIDOC
## discrete_grouped_gemm_swiglu_wrapper_sm100

### Description
Executes a discrete grouped GEMM operation with SwiGLU activation using the provided input tensors and configuration parameters.

### Parameters
- **a_tensor** (Tensor) - Input tensor A.
- **b_ptrs** (Tensor) - Pointers to per-expert B tensors.
- **sfa_tensor** (Tensor) - Scale factor tensor for A.
- **sfb_ptrs** (Tensor) - Pointers to per-expert scale factor tensors for B.
- **padded_offsets** (Tensor) - Padded offsets for grouping.
- **alpha_tensor** (Tensor) - Alpha scaling tensor.
- **n** (int) - Logical full N dimension before GLU split.
- **b_dtype** (dtype) - Data type of per-expert B tensors.
- **norm_const_tensor** (Tensor) - Normalization constant tensor (required if SFD outputs enabled).
- **prob_tensor** (Tensor) - Optional probability tensor.
- **acc_dtype** (dtype) - Accumulation data type.
- **c_dtype** (dtype) - Data type for intermediate GEMM result.
- **d_dtype** (dtype) - Data type for row-quantized GLU output.
- **cd_major** (str) - Layout major for C/D tensors.
- **mma_tiler_mn** (tuple) - MMA tiler dimensions.
- **cluster_shape_mn** (tuple) - Cluster shape dimensions.
- **sf_vec_size** (int) - Scale factor vector size.
- **vector_f32** (bool) - Whether to use vector F32.
- **m_aligned** (int) - M alignment value.
- **discrete_col_sfd** (bool) - Whether to use discrete column SFD.
- **act_func** (str) - Activation function (e.g., "swiglu", "geglu").
- **b_major** (str) - Major layout for B tensors.
- **current_stream** (CUstream) - CUDA stream for execution.

### Returns
- **outputs** (dict/tuple) - Contains 'c_tensor', 'd_tensor', 'd_col_tensor', 'amax_tensor', 'sfd_row_tensor', and 'sfd_col_tensor'.
```

--------------------------------

### create_execution_plan

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/utilities/custom-execution-plan.html

Creates an execution plan using a specific engine and a set of knobs.

```APIDOC
## create_execution_plan

### Description
Creates an execution plan with a particular engine and configuration knobs.

### Signature
`error_t create_execution_plan(int64_t const engine_id, std::unordered_map<KnobType_t, int64_t> const &knobs);`

### Parameters
- **engine_id** (int64_t const) - The engine index.
- **knobs** (std::unordered_map<KnobType_t, int64_t> const &) - The configuration knobs for the plan.
```

--------------------------------

### Class API Usage

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.html

Demonstrates the class-based API for Grouped GEMM DGLU, requiring explicit compilation before execution.

```python
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,
)
```

```python
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,
)
```

--------------------------------

### Initialize and Set Kernel Cache

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/utilities/dynamic-kernel-cache.html

Creates a kernel cache instance and associates it with a dynamic shape graph to enable re-use of compiled kernels.

```cpp
auto kernel_cache = std::make_shared<cudnn_frontend::KernelCache>();
```

```cpp
graph.set_kernel_cache(kernel_cache)
```

--------------------------------

### Access NSA Components

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/nsa.html

Import the NSA namespace to access available attention components and their respective wrappers.

```python
from cudnn import NSA

# Access individual components
NSA.SelectionAttention         # Class API
NSA.selection_attention_wrapper  # High-level wrapper

NSA.CompressionAttention
NSA.compression_attention_wrapper

NSA.SlidingWindowAttention
NSA.sliding_window_attention_wrapper

NSA.TopKReduction
NSA.topk_reduction_wrapper
```

--------------------------------

### Execute Indexer Top-K

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/dsa.html

Select candidate KV indices from indexer scores using a radix top-K kernel. Requires SM100+ hardware.

```python
result = DSA.indexer_top_k_wrapper(
    scores.reshape(-1, scores.shape[-1]),
    seq_lens, top_k=512,
)
indices, values = result["indices"], result["values"]
```

--------------------------------

### DiscreteGroupedGemmDswigluSm100.__init__

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_dswiglu.html

Constructor for the DiscreteGroupedGemmDswigluSm100 class. Initializes the operator with sample tensors and configuration parameters.

```APIDOC
## DiscreteGroupedGemmDswigluSm100.__init__

### Description
Initializes the DiscreteGroupedGemmDswigluSm100 operator. Requires sample tensors to define the shapes and types for the GEMM operation.

### Parameters
- **sample_a**, **num_experts**, **b_shape**, **b_dtype**, **sample_c**, **sample_d_row**, **sample_d_col**, **sample_sfa**, **sample_padded_offsets**, **sample_alpha**, **sample_beta**, **sample_prob**, **sample_dprob**, **sample_sfd_row**, **sample_sfd_col**, **sample_amax**, **sample_norm_const** (Various) - Required - Input/Output tensor samples defining the operation structure. Note: `sample_sfd_row`, `sample_sfd_col`, and `sample_norm_const` must be all None or all provided.
```

--------------------------------

### GemmSwigluSm100 Constructor

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/gemm_swiglu.html

Initializes the GemmSwigluSm100 instance with sample tensors for standard or quantized mode operations.

```APIDOC
## GemmSwigluSm100(sample_a, sample_b, sample_ab12, sample_c, sample_sfa=None, sample_sfb=None, sample_sfc=None, sample_amax=None, sample_norm_const=None)

### Description
Initializes the GEMM SwiGLU operator for SM100 compute capability. The constructor accepts sample tensors to define the layout and configuration for subsequent execution.

### Parameters
- **sample_a** (tensor) - Required - Input tensor A
- **sample_b** (tensor) - Required - Input tensor B
- **sample_ab12** (tensor) - Required - Intermediate GEMM output tensor
- **sample_c** (tensor) - Required - SwiGLU output tensor
- **sample_sfa** (tensor) - Optional - Scale factor tensor A (quantized mode)
- **sample_sfb** (tensor) - Optional - Scale factor tensor B (quantized mode)
- **sample_sfc** (tensor) - Optional - Scale factor tensor C (quantized mode)
- **sample_amax** (tensor) - Optional - Max-abs output tensor (quantized mode)
- **sample_norm_const** (tensor) - Optional - Normalization constant tensor (quantized mode)
```

--------------------------------

### Check Graph Support

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/developer/overview.html

Verifies that the graph can be successfully executed using the queried plans.

```cpp
cudnn_frontend::error_t
cudnn_frontend::graph::Graph::check_support(cudnnHandle_t h);
```

--------------------------------

### Execute GEMM SwiGLU with Class API

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/gemm_swiglu.html

Use the class-based API for more control, requiring explicit support checks and compilation before execution.

```python
gemm = GemmSwigluSm100(
    sample_a,
    sample_b,
    sample_ab12,
    sample_c,
    alpha=1.0,
    acc_dtype=torch.float32,
    mma_tiler_mn=(128, 128),
    cluster_shape_mn=None,
)
assert gemm.check_support()
gemm.compile()
gemm.execute(
    a_tensor,
    b_tensor,
    ab12_tensor,
    c_tensor,
    alpha=1.0,
    current_stream=None,
)
```

```python
gemm = GemmSwigluSm100(
    sample_a,
    sample_b,
    sample_ab12,
    sample_c,
    alpha=1.0,
    acc_dtype=torch.float32,
    mma_tiler_mn=(128, 128),
    cluster_shape_mn=None,
    # Quantization parameters
    sample_sfa=sample_sfa,
    sample_sfb=sample_sfb,
    sample_amax=sample_amax,          # Required for fp4 inputs with bf16 output
    sample_sfc=sample_sfc,            # Required when c_dtype is fp8
    sample_norm_const=sample_norm_const,  # Required when c_dtype is fp8
    sf_vec_size=16,
    vector_f32=False,
    ab12_stages=4,
)
assert gemm.check_support()
gemm.compile()
gemm.execute(
    a_tensor,
    b_tensor,
    ab12_tensor,
    c_tensor,
    sfa_tensor=sfa_tensor,
    sfb_tensor=sfb_tensor,
    amax_tensor=amax_tensor,
    sfc_tensor=sfc_tensor,
    norm_const_tensor=norm_const_tensor,
    alpha=1.0,
    current_stream=None,
)
```

--------------------------------

### Create a public API wrapper

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/utilities/adding_torch_custom_ops.html

Expose a user-friendly function that follows standard PyTorch conventions for the custom operation.

```python
def my_op(x, w, eps=1e-5, bias=None):
    """cuDNN-accelerated my_op. Matches torch.nn.functional.my_op API."""
    # input validation
    # ...
    y, _aux = torch.ops.cudnn.my_op(x, w, eps, bias=bias)
    return y
```

--------------------------------

### Query Autotune Workspace Size

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/developer/overview.html

Retrieve the workspace size required to run autotuning across all plans.

```cpp
get_autotune_workspace_size() const
```

--------------------------------

### Build Execution Plans

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/developer/overview.html

Constructs execution plans using either a policy-based approach or by targeting specific plan indices.

```cpp
cudnn_frontend::error_t
cudnn_frontend::graph::Graph::build_plan(
    cudnnHandle_t const &handle,
    cudnn_frontend::BuildPlanPolicy_t const policy,
    bool const do_multithreaded_builds
);
```

```cpp
cudnn_frontend::error_t
cudnn_frontend::Graph::build_plan_at_index(
    cudnnHandle_t const &handle,
    int64_t plan_index
);
```

--------------------------------

### NSA.CompressionAttention Class

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/nsa.html

A class-based API for configuring, compiling, and executing compression attention operations.

```APIDOC
## NSA.CompressionAttention

### Description
Class for managing the lifecycle of a compression attention operation, including support checking, compilation, and execution.

### Methods
- **check_support()** - Validates if the current configuration is supported.
- **compile()** - Compiles the attention kernel.
- **execute(...)** - Executes the attention operation with provided tensors and stream.
```

--------------------------------

### GroupedGemmQuantSm100.__init__

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/grouped_gemm_quant.html

Constructor for the GroupedGemmQuantSm100 class. Initializes the operator with sample tensors defining the expected layout and shapes.

```APIDOC
## GroupedGemmQuantSm100.__init__

### Description
Initializes the GroupedGemmQuantSm100 operator. The constructor requires sample tensors for all inputs and outputs to define the expected layout and shapes for subsequent execution.

### Parameters
- **sample_a**, **sample_b**, **sample_d**, **sample_sfa**, **sample_sfb**, **sample_padded_offsets**, **sample_alpha**, **sample_d_col**, **sample_sfd_row**, **sample_sfd_col**, **sample_amax**, **sample_norm_const**, **sample_prob** (Tensor) - Required - Sample tensors defining the layout. Note: `sample_sfd_row`, `sample_sfd_col`, and `sample_norm_const` must be all None or all provided.
```

--------------------------------

### Execute Indexer Forward

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/dsa.html

Compute dense indexer scores with a ratio-causal mask. Requires SM100+ hardware.

```python
result = DSA.indexer_forward_wrapper(
    q, k, w, ratio=4, q_causal_offsets=q_causal_offsets,
)
scores = result["scores"]
```

--------------------------------

### DiscreteGroupedGemmSwigluSm100 Constructor

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.html

Initializes the DiscreteGroupedGemmSwigluSm100 class for performing grouped GEMM SwiGLU operations.

```APIDOC
## DiscreteGroupedGemmSwigluSm100.__init__

### Description
Initializes the class instance for the Discrete Grouped GEMM SwiGLU operation on SM100 hardware.

### Parameters
- **Parameters** (dict) - Configuration parameters for the operation, including tiling, cluster, and data type constraints.
```

--------------------------------

### Execute GEMM SwiGLU with High-level Wrapper

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/gemm_swiglu.html

Use the high-level wrapper for standard or quantized GEMM SwiGLU operations. Quantized mode is triggered by providing scale factor tensors.

```python
result = gemm_swiglu_wrapper_sm100(
    a_tensor,
    b_tensor,
    alpha=1.0,
    c_major="m",
    ab12_dtype=torch.float32,
    c_dtype=torch.float16,
    acc_dtype=torch.float32,
    mma_tiler_mn=(128, 128),
    cluster_shape_mn=(1, 1),
    stream=None,
)
ab12, c, sfc, amax = result
# Note: sfc and amax are always None in standard mode
# Key access: result["ab12_tensor"], result["c_tensor"]
```

```python
result = gemm_swiglu_wrapper_sm100(
    a_tensor,
    b_tensor,
    alpha=1.0,
    c_major="m",
    ab12_dtype=torch.bfloat16,
    c_dtype=torch.bfloat16,
    acc_dtype=torch.float32,
    mma_tiler_mn=(128, 128),
    cluster_shape_mn=(1, 1),
    # Quantization parameters
    sfa_tensor=sfa_tensor,
    sfb_tensor=sfb_tensor,
    norm_const_tensor=norm_const_tensor,  # Required when c_dtype is fp8
    sf_vec_size=16,
    vector_f32=False,
    ab12_stages=4,
    stream=None,
)
ab12, c, sfc, amax = result
# Key access: result["ab12_tensor"], result["c_tensor"], result["sfc_tensor"], result["amax_tensor"]
```

--------------------------------

### DiscreteGroupedGemmSwigluSm100 Constructor

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.html

Initializes the DiscreteGroupedGemmSwigluSm100 class with sample tensors and configuration parameters.

```APIDOC
## DiscreteGroupedGemmSwigluSm100.__init__

### Description
Initializes the operator with sample tensors and configuration parameters to define the grouped GEMM operation.

### Parameters
- **sample_a** (tensor) - Input tensor A
- **num_experts** (int) - Number of experts
- **b_shape** (tuple) - Logical (N, K) shape for one expert
- **b_dtype** (torch.dtype) - Dtype of expert B tensors
- **sample_c** (tensor) - Sample intermediate C tensor
- **sample_d** (tensor) - Sample output D tensor
- **sample_sfa** (tensor) - Sample scale factor A
- **sample_padded_offsets** (tensor) - Sample padded offsets
- **sample_alpha** (tensor) - Sample alpha tensor
- **sample_d_col** (tensor) - Sample column-quantized output
- **sample_sfd_row** (tensor) - Sample row scale factors
- **sample_sfd_col** (tensor) - Sample column scale factors
- **sample_amax** (tensor) - Sample amax tensor
- **sample_norm_const** (tensor) - Sample normalization constant
- **sample_prob** (tensor) - Sample probability tensor
```

--------------------------------

### Discrete-weight wrapper usage

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.html

Executes the grouped GEMM DSReLU operation specifically for discrete-weight configurations.

```python
result = grouped_gemm_dsrelu_wrapper_sm100(
    a_tensor=a,
    c_tensor=c,
    sfa_tensor=sfa,
    padded_offsets=padded_offsets,
    alpha_tensor=alpha,
    prob_tensor=prob,
    b_ptrs=b_ptrs,          # int64 device tensor of per-expert B pointers
    sfb_ptrs=sfb_ptrs,      # int64 device tensor of per-expert SFB pointers
    n=N,
    b_dtype=torch.float4_e2m1fn_x2,
    b_major="k",
    d_dtype=torch.bfloat16,
)
```

--------------------------------

### grouped_gemm_dglu_wrapper_sm100

Source: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.html

A high-level functional wrapper for executing Grouped GEMM DGLU operations on SM100 hardware. It supports both dense and discrete weight modes.

```APIDOC
## grouped_gemm_dglu_wrapper_sm100

### Description
Executes a Grouped GEMM DGLU operation. This function accepts tensors for inputs and weights, and returns a dictionary or tuple containing the output tensors.

### Usage
```python
from cudnn import grouped_gemm_dglu_wrapper_sm100

# Dense mode example
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, b_tensor=b, sfb_tensor=sfb, 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", current_stream=stream
)
```
```
Tokens:
Raw