### NSA Pipeline Initialization Example

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

This example demonstrates the initial setup for a full NSA pipeline, including defining configuration parameters such as batch size, sequence length, head dimensions, and attention mechanism specifics like k_value, block size, and stride. It also sets the device and data type for the operations.

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

```

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

### Low-level Graph API: moe_grouped_matmul Example

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

A complete example demonstrating the setup and execution of moe_grouped_matmul using the low-level cuDNN graph API. It includes tensor definitions and graph building steps.

```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()


```

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

### Install Python Dependencies

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

Installs the necessary dependencies to run the Python cuDNN samples. Ensure you have pip installed.

```bash
pip install -r requirements.txt

```

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

### Grouped GEMM + SwiGLU Wrapper Usage (Python)

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

Example of using the high-level Python wrapper for the Grouped GEMM + SwiGLU fusion. Ensure CUDA stream is correctly passed.

```python
from cudnn import grouped_gemm_swiglu_wrapper_sm100
from cuda.bindings import driver as cuda

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

outputs = grouped_gemm_swiglu_wrapper_sm100(
    a_tensor=a,
    b_tensor=b,
    sfa_tensor=sfa,
    sfb_tensor=sfb,
    padded_offsets=padded_offsets,
    alpha_tensor=alpha,
    norm_const_tensor=norm_const,  # Required when SFD outputs are enabled (FP8 inputs)
    prob_tensor=prob,
    acc_dtype=torch.float32,
    c_dtype=torch.bfloat16,
    d_dtype=torch.bfloat16,
    cd_major="n",
    mma_tiler_mn=(256, 256),
    cluster_shape_mn=(2, 1),
    sf_vec_size=32,
    vector_f32=False,
    m_aligned=256,
    discrete_col_sfd=False,
    current_stream=stream,
)

```

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

### Python SDPA Forward Pass Example

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

Example demonstrating how to set up and execute a forward pass for Scaled Dot Product Attention using cuDNN in Python. It includes tensor creation, graph building, and execution.

```python
import cudnn
import torch
import math

# Create graph
graph = cudnn.pygraph(
    io_data_type=cudnn.data_type.HALF,
    intermediate_data_type=cudnn.data_type.FLOAT,
    compute_data_type=cudnn.data_type.FLOAT,
)

# Create tensor descriptors
q = graph.tensor_like(q_gpu)
k = graph.tensor_like(k_gpu)
v = graph.tensor_like(v_gpu)

# Forward pass with causal masking
o, stats = graph.sdpa(
    name="sdpa",
    q=q,
    k=k,
    v=v,
    attn_scale=1.0 / math.sqrt(d),
    generate_stats=True,                          # For training
    diagonal_band_right_bound=0,                  # Causal mask
    diagonal_alignment=cudnn.diagonal_alignment.TOP_LEFT,
)

o.set_output(True).set_dim(shape_o).set_stride(stride_o)
stats.set_output(True).set_data_type(cudnn.data_type.FLOAT)

# Build and execute
graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])


```

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

### Get Execution Plans

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

Internally queries heuristics for engine configurations based on the provided heuristics modes.

```cpp
cudnn_frontend::error_t
cudnn_frontend::graph::Graph::get_execution_plans(std::vector<heur_mode_t>)
```

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

### Wrapper API Example for Grouped GEMM Wgrad

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

Demonstrates how to use the `grouped_gemm_wgrad_wrapper_sm100` function to compute grouped weight gradients. This wrapper supports 'dense' output mode and specifies the desired weight gradient data type.

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

wgrad_tensor = result["wgrad_tensor"]

```

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

### Install cuDNN Frontend with cutedsl dependency

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

Install the nvidia-cudnn-frontend package with the optional 'cutedsl' dependency using pip. This is required for operations like GEMM + Amax and GEMM + SwiGLU.

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

```

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

### Example Usage of Scaled Dot-Product Attention

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

Demonstrates how to use the scaled_dot_product_attention function with sample tensors. Includes forward pass and notes on backward pass handling by autograd.

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

```

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

### Class API Example for Grouped GEMM Wgrad

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

Shows how to instantiate and use the `cudnn.GroupedGemmWgradSm100` class for grouped weight gradient computation. This involves checking support, compiling the operation, and then executing it with input and output tensors.

```python
import cudnn
import torch

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

```

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

### Initialize and Execute Compression Attention

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

Shows how to instantiate the CompressionAttention class with sample tensors and configuration, then check support, compile, and execute the attention operation.

```python
from cudnn import NSA

comp_attn = NSA.CompressionAttention(
    sample_q=q,
    sample_k=k,
    sample_v=v,
    sample_o=o,
    sample_lse=lse,  # Optional, set to None for inference
    sample_cum_seqlen_q=cum_seqlen_q,
    sample_cum_seqlen_k=cum_seqlen_k,
    qk_acc_dtype=torch.float32,
    pv_acc_dtype=torch.float32,
    mma_tiler_mn=(128, 128),
    is_persistent=False,
    scale_q=1.0,
    scale_k=1.0,
    scale_v=1.0,
    inv_scale_o=1.0,
    scale_softmax=None,
)
assert comp_attn.check_support()
comp_attn.compile()
comp_attn.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,
    current_stream=stream,
)
```

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

### Initialize and Execute Selection Attention

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

Demonstrates how to initialize the SelectionAttention class, check for support, compile the kernel, and execute the attention mechanism with provided tensors and stream.

```python
from cudnn import NSA
from cuda.bindings import driver as cuda

selection_attention = NSA.SelectionAttention(
    sample_q=q,
    sample_k=k,
    sample_v=v,
    sample_o=o,
    sample_l=l,
    sample_m=m,
    sample_block_indices=block_indices,
    sample_block_counts=block_counts,
    sample_cum_seqlen_q=cum_seqlen_q,
    sample_cum_seqlen_k=cum_seqlen_k,
    max_s_q=1024,
    max_s_k=1024,
    acc_dtype=torch.float32,
    block_size=64,
    scale_softmax=None,
)
assert selection_attention.check_support()
selection_attention.compile()
selection_attention.execute(
    q_tensor=q,
    k_tensor=k,
    v_tensor=v,
    o_tensor=o,
    l_tensor=l,
    m_tensor=m,
    block_indices_tensor=block_indices,
    block_counts_tensor=block_counts,
    cum_seqlen_q_tensor=cum_seqlen_q,
    cum_seqlen_k_tensor=cum_seqlen_k,
    scale_softmax=None,
    current_stream=stream,
)
```

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

### CausalConv1d Example with Bias and Silu Activation

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

Example demonstrating the usage of cudnn.ops.causal_conv1d with a bias tensor and 'silu' activation, including backward pass for gradient computation.

```python
import torch
import cudnn

B, D, L, K = 2, 768, 4096, 4

x = torch.randn(B, D, L, dtype=torch.bfloat16, device="cuda", requires_grad=True)
w = torch.randn(D, K, dtype=torch.bfloat16, device="cuda", requires_grad=True)
b = torch.randn(D, dtype=torch.bfloat16, device="cuda", requires_grad=True)

y = cudnn.ops.causal_conv1d(x, w, bias=b, activation="silu")

loss = y.sum()
loss.backward()  # x.grad, w.grad, b.grad are populated


```

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

### Initialize Device Properties

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

Demonstrates the different methods to initialize a device property descriptor. Choose the method that best suits your available information, such as a cuDNN handle, a specific device ID, or a serialized buffer.

```c++
set_handle(cudnnHandle_t handle);  // initialize from a cuDNN handle
```

```c++
set_device_id(int32_t device_id);  // initialize from a specific device
```

```c++
deserialize(const std::vector<uint8_t>& serialized_buf);  // deserialize from json
```

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

### Instantiate and Execute GroupedGemmQuantSm100 API

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

Demonstrates how to instantiate the GroupedGemmQuantSm100 API with various sample tensors and configuration options, followed by compilation and execution. Ensure all required sample tensors and configuration parameters are provided.

```python
from cudnn import GroupedGemmQuantSm100
from cuda.bindings import driver as cuda

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

api = GroupedGemmQuantSm100(
    sample_a=a,
    sample_b=b,
    sample_d=d,
    sample_sfa=sfa,
    sample_sfb=sfb,
    sample_padded_offsets=padded_offsets,
    sample_alpha=alpha,
    sample_d_col=d_col,
    # Optional quantization outputs
    sample_sfd_row=sfd_row,        # Required when SFD outputs are enabled
    sample_sfd_col=sfd_col,        # Required when SFD outputs are enabled
    sample_amax=amax,              # Required for bf16 output with FP4 input
    sample_norm_const=norm_const,  # Required when SFD outputs are enabled
    sample_prob=prob,              # Per-row gating probabilities (required)
    # Configuration
    acc_dtype=torch.float32,
    mma_tiler_mn=(256, 256),
    cluster_shape_mn=(2, 1),
    sf_vec_size=16,
    vector_f32=False,
    m_aligned=256,
    discrete_col_sfd=False,
)
assert api.check_support()
api.compile()
api.execute(
    a_tensor=a,
    b_tensor=b,
    d_tensor=d,
    sfa_tensor=sfa,
    sfb_tensor=sfb,
    padded_offsets=padded_offsets,
    alpha_tensor=alpha,
    d_col_tensor=d_col,
    sfd_row_tensor=sfd_row,
    sfd_col_tensor=sfd_col,
    amax_tensor=amax,
    norm_const_tensor=norm_const,
    prob_tensor=prob,
    current_stream=stream,
)
```

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

### Getting the Execution Plan Count

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

Returns the number of execution plans returned by cuDNN heuristics. Each plan gets an index from 0 to #plans-1, with 0 having top priority.

```APIDOC
## Getting the Execution Plan Count

This method returns the number of execution plans returned by cuDNN heuristics. Each plan gets an index from `0` to `#plans-1`, with `0` having top priority.

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

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

### Class API (Standard Mode)

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

Instantiate and use the `GemmSwigluSm100` class for standard mode GEMM + SwiGLU fusion. Requires sample tensors for initialization and provides methods to check support, compile, and execute the operation.

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

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

### DeviceProperties Descriptor Creation and Initialization

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

Demonstrates how to create a DeviceProperties descriptor and initialize it using different methods: from a cuDNN handle, a device ID, or by deserializing from a buffer.

```APIDOC
## DeviceProperties Descriptor

### Description
Creates a descriptor for cuDNN device properties, which can be serialized and used to query cuDNN heuristics or create execution plans without the physical device being available.

### API

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

**Initialization Methods:**

1.  **Initialize from a cuDNN handle:**
    ```cpp
    device_prop->set_handle(cudnnHandle_t handle);
    ```

2.  **Initialize from a specific device ID:**
    ```cpp
    device_prop->set_device_id(int32_t device_id);
    ```

3.  **Deserialize from a buffer (e.g., JSON):
    ```cpp
    device_prop->deserialize(const std::vector<uint8_t>& serialized_buf);
    ```

### Usage Example

Set device properties for a graph:

```cpp
graph.set_device_properties(device_prop);
```
```

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

### Build cuDNN C++ Samples with CMake

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

Use this command to configure and build the cuDNN C++ samples. Ensure CUDA and cuDNN paths are correctly set.

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

```

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

### Class API

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

A class-based API for configuring and executing the GEMM + sReLU fusion. It allows for detailed setup and execution of the operation.

```APIDOC
## Class API

### Description
A class-based API for configuring and executing the GEMM + sReLU fusion. It allows for detailed setup and execution of the operation.

### Method
```python
from cudnn import GemmSreluSm100

op = GemmSreluSm100(
    sample_a=a,
    sample_b=b,
    sample_c=c,
    sample_d=d,
    sample_sfa=sfa,
    sample_sfb=sfb,
    sample_prob=prob,
    sample_sfd=sfd,
    sample_amax=amax,
    sample_norm_const=norm_const,
    alpha=1.0,
    acc_dtype=torch.float32,
    mma_tiler_mn=(256, 256),
    cluster_shape_mn=(2, 1),
    sf_vec_size=16,
    vector_f32=False,
)
assert op.check_support()
op.compile()
op.execute(
    a_tensor=a,
    b_tensor=b,
    c_tensor=c,
    d_tensor=d,
    sfa_tensor=sfa,
    sfb_tensor=sfb,
    prob_tensor=prob,
    sfd_tensor=sfd,
    amax_tensor=amax,
    norm_const_tensor=norm_const,
    current_stream=None,
)
```

### Parameters for Initialization
* **sample_a**: Sample tensor for input A.
* **sample_b**: Sample tensor for input B.
* **sample_c**: Sample tensor for output C.
* **sample_d**: Sample tensor for output D.
* **sample_sfa**: Sample tensor for scale factor A.
* **sample_sfb**: Sample tensor for scale factor B.
* **sample_prob**: Sample tensor for probability.
* **sample_sfd**: Sample tensor for output scale factor D.
* **sample_amax**: Sample tensor for Amax.
* **sample_norm_const**: Sample tensor for normalization constant.
* **alpha**: Scaling factor for the GEMM operation (default: 1.0).
* **acc_dtype**: Data type for accumulation (default: torch.float32).
* **mma_tiler_mn**: MMA tiler shape (default: (256, 256)).
* **cluster_shape_mn**: Cluster shape (default: (2, 1)).
* **sf_vec_size**: Scale factor vector size (default: 16).
* **vector_f32**: Whether to use vector FP32 (default: False).

### Methods
* **check_support()**: Checks if the current hardware supports the operation.
* **compile()**: Compiles the kernel for the operation.
* **execute(...)**: Executes the GEMM + sReLU fusion with the provided tensors and stream.
```

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

### Get Engine Count

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

Retrieves the total number of available engines. This function is used to determine how many execution engines can be utilized.

```cpp
inline error_t
get_engine_count(int64_t &count);
```

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

### Instantiate and Execute GemmSwigluSm100 in Quantized Mode

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

This snippet demonstrates how to instantiate the GemmSwigluSm100 class with quantization parameters and then execute the operation. Ensure all required quantization-specific tensors (e.g., sample_sfa, sample_sfb, sample_amax) are provided based on the input and output dtypes.

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

```

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

### Using the High-level Wrapper for Discrete Grouped GEMM SWIGLU

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

Demonstrates how to call the discrete_grouped_gemm_swiglu_wrapper_sm100 function with various input tensors and configuration parameters. Shows how to access the output tensors via dictionary keys or tuple unpacking.

```python
from cudnn import discrete_grouped_gemm_swiglu_wrapper_sm100
from cuda.bindings import driver as cuda

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

b_ptrs = torch.tensor([b.data_ptr() for b in b_list], dtype=torch.int64, device="cuda")
sfb_ptrs = torch.tensor([sfb.data_ptr() for sfb in sfb_list], dtype=torch.int64, device="cuda")

outputs = discrete_grouped_gemm_swiglu_wrapper_sm100(
    a_tensor=a_tensor,
    b_ptrs=b_ptrs,
    sfa_tensor=sfa_tensor,
    sfb_ptrs=sfb_ptrs,
    padded_offsets=padded_offsets,
    alpha_tensor=alpha_tensor,
    n=n,                              # logical full N before GLU split
    b_dtype=b_dtype,                  # dtype of per-expert B tensors
    norm_const_tensor=norm_const,     # required when SFD outputs are enabled
    prob_tensor=prob_tensor,          # optional
    acc_dtype=torch.float32,
    c_dtype=torch.bfloat16,
    d_dtype=torch.bfloat16,
    cd_major="n",
    mma_tiler_mn=(256, 256),
    cluster_shape_mn=(2, 1),
    sf_vec_size=32,
    vector_f32=False,
    m_aligned=256,
    discrete_col_sfd=False,
    act_func="swiglu",                # or "geglu"
    b_major="k",                      # or "n" (fp8 only)
    current_stream=stream,
)

# dictionary access:
c = outputs["c_tensor"]               # intermediate GEMM result
d = outputs["d_tensor"]               # row-quantized GLU output
d_col = outputs["d_col_tensor"]       # column-quantized GLU output
amax = outputs["amax_tensor"]         # per-group amax (when d_dtype is bf16/fp16)
sfd_row = outputs["sfd_row_tensor"]   # row scale factors (when enabled)
sfd_col = outputs["sfd_col_tensor"]   # column scale factors (when enabled)

# or tuple unpacking:
c, d, d_col, amax, sfd_row, sfd_col = outputs

```

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

### API Usage with Python and cuDNN

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

Example demonstrating how to use the cuDNN graph API to build the Fused RMSNorm + SiLU pattern with Python.

```python
import cudnn
import torch

C, num_tokens = 512, 24960
eps = 1e-5

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

X = graph.tensor(
    dim=[num_tokens, C, 1, 1],
    stride=[C, 1, 1, 1],
    data_type=cudnn.data_type.BFLOAT16,
)
scale = graph.tensor(
    dim=[1, C, 1, 1],
    stride=[C, 1, 1, 1],
    data_type=cudnn.data_type.BFLOAT16,
)
epsilon = graph.tensor(
    dim=[1, 1, 1, 1],
    stride=[1, 1, 1, 1],
    data_type=cudnn.data_type.FLOAT,
    is_pass_by_value=True,
)

# Build the RMSNorm → SiLU fusion pattern
Y = graph.rmsnorm(
    norm_forward_phase=cudnn.norm_forward_phase.INFERENCE,
    input=X, scale=scale, epsilon=epsilon,
)[0]
Y.set_dim([num_tokens, C, 1, 1]).set_stride([C, 1, 1, 1])
Y.set_data_type(cudnn.data_type.BFLOAT16)

Z = graph.swish(input=Y, swish_beta=1.0)
Z.set_output(True).set_data_type(cudnn.data_type.BFLOAT16)

# Build with OPENSOURCE heuristic mode
graph.build([cudnn.heur_mode.OPENSOURCE])

```

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

### Get Knobs for Engine

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

Fetches the list of supported knobs for a given engine. This is useful for understanding the configurable parameters of an engine before creating a plan.

```cpp
inline error_t
get_knobs_for_engine(int64_t const engine, std::vector<Knob> &);

```

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

### 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);
```

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

### Registering Custom Autograd Function

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

Defines the setup and backward functions for a custom cuDNN operation and registers it with PyTorch's autograd system.

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

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

### Get Execution Plan Count

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

Returns the number of execution plans generated by cuDNN heuristics. Plans are indexed from 0, with 0 having the highest priority.

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

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

### Initialize Host Tensors and Copy to Device

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

Initializes host vectors A_host and B_host with random values, then casts them to nv_bfloat16 and copies them to the device tensors A_dev and B_dev.

```cpp
// Host copies kept for CPU reference computation.
    std::vector<float> A_host(b * m * k);
    std::vector<float> B_host(1 * k * n);

    // Initialize the CPU reference tensors with random values, then make a copy to the GPU.
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<float> dis(0.0f, 1.0f);
    for (auto& e : A_host) e = dis(gen);
    for (auto& e : B_host) e = dis(gen);

    for (int64_t i = 0; i < b * m * k; ++i) {
        A_dev[i] = __float2bfloat16(A_host[i]);
    }
    for (int64_t i = 0; i < 1 * k * n; ++i) {
        B_dev[i] = __float2bfloat16(B_host[i]);
    }
    cudaDeviceSynchronize();
```

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

### Matmul Attributes Setters C++ API

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

These are examples of setters for the Matmul_attributes structure in the C++ API. They allow configuring the operation's name and compute data type.

```cpp
Matmul_attributes& set_name(std::string const&)

Matmul_attributes& set_compute_data_type(DataType_t value)
```

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

### Combine NSA Outputs

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

Combines the outputs from different NSA attention mechanisms. This example shows a placeholder for a weighted combination, which in practice might involve learned gating mechanisms.

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

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

### GroupedGemmQuantSm100 Initialization and Execution

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

This snippet demonstrates how to initialize, compile, and execute the GroupedGemmQuantSm100 API. It covers the necessary parameters for both initialization and execution, including optional quantization outputs.

```APIDOC
## GroupedGemmQuantSm100 API Usage

### Description
This section outlines the process of using the `GroupedGemmQuantSm100` class for quantized grouped matrix multiplication. It includes initialization with sample tensors and configuration, followed by compilation and execution.

### Initialization
Initialize the `GroupedGemmQuantSm100` object with input tensors and configuration parameters.

```python
from cudnn import GroupedGemmQuantSm100
from cuda.bindings import driver as cuda
import torch

# Assuming a, b, d, sfa, sfb, padded_offsets, alpha, norm_const, prob are defined tensors
# and stream is a valid CUstream
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

api = GroupedGemmQuantSm100(
    sample_a=a,
    sample_b=b,
    sample_d=d,
    sample_sfa=sfa,
    sample_sfb=sfb,
    sample_padded_offsets=padded_offsets,
    sample_alpha=alpha,
    sample_d_col=d_col,
    # Optional quantization outputs
    sample_sfd_row=sfd_row,        # Required when SFD outputs are enabled
    sample_sfd_col=sfd_col,        # Required when SFD outputs are enabled
    sample_amax=amax,              # Required for bf16 output with FP4 input
    sample_norm_const=norm_const,  # Required when SFD outputs are enabled
    sample_prob=prob,              # Per-row gating probabilities (required)
    # Configuration
    acc_dtype=torch.float32,
    mma_tiler_mn=(256, 256),
    cluster_shape_mn=(2, 1),
    sf_vec_size=16,
    vector_f32=False,
    m_aligned=256,
    discrete_col_sfd=False,
)

# Check support and compile the API
assert api.check_support()
api.compile()

### Execution
Execute the compiled API with the actual input tensors.

```python
api.execute(
    a_tensor=a,
    b_tensor=b,
    d_tensor=d,
    sfa_tensor=sfa,
    sfb_tensor=sfb,
    padded_offsets=padded_offsets,
    alpha_tensor=alpha,
    d_col_tensor=d_col,
    sfd_row_tensor=sfd_row,
    sfd_col_tensor=sfd_col,
    amax_tensor=amax,
    norm_const_tensor=norm_const,
    prob_tensor=prob,
    current_stream=stream,
)
```

### Parameters

#### Initialization Parameters

- **sample_a** (Tensor) - Sample tensor for input A.
- **sample_b** (Tensor) - Sample tensor for input B.
- **sample_d** (Tensor) - Sample tensor for output D.
- **sample_sfa** (Tensor) - Sample tensor for scale factor A.
- **sample_sfb** (Tensor) - Sample tensor for scale factor B.
- **sample_padded_offsets** (Tensor) - Sample tensor for padded offsets.
- **sample_alpha** (Tensor) - Sample tensor for alpha.
- **sample_d_col** (Tensor or None) - Sample tensor for column-wise output D (None for bf16/fp16/fp32 outputs).
- **sample_sfd_row** (Tensor, optional) - Sample tensor for row scale factors (required when SFD outputs are enabled).
- **sample_sfd_col** (Tensor, optional) - Sample tensor for column scale factors (required when SFD outputs are enabled).
- **sample_amax** (Tensor, optional) - Sample tensor for per-group amax (required for bf16 output with FP4 input).
- **sample_norm_const** (Tensor, optional) - Sample tensor for normalization constant (required when SFD outputs are enabled).
- **sample_prob** (Tensor) - Sample tensor for per-row gating probabilities (required).
- **acc_dtype** (torch.dtype) - Accumulation data type.
- **mma_tiler_mn** (tuple) - MMA tiler shape (M, N).
- **cluster_shape_mn** (tuple) - Cluster shape (M, N).
- **sf_vec_size** (int) - Vector size for scale factors.
- **vector_f32** (bool) - Whether to use vector FP32.
- **m_aligned** (int) - Alignment for dimension M.
- **discrete_col_sfd** (bool) - Whether to use discrete column SFD.

#### Execution Parameters

- **a_tensor** (Tensor) - Input tensor A.
- **b_tensor** (Tensor) - Input tensor B.
- **d_tensor** (Tensor) - Output tensor D.
- **sfa_tensor** (Tensor) - Scale factor A tensor.
- **sfb_tensor** (Tensor) - Scale factor B tensor.
- **padded_offsets** (Tensor) - Padded offsets tensor.
- **alpha_tensor** (Tensor) - Alpha tensor.
- **d_col_tensor** (Tensor or None) - Column-wise output D tensor.
- **sfd_row_tensor** (Tensor, optional) - Row scale factors tensor.
- **sfd_col_tensor** (Tensor, optional) - Column scale factors tensor.
- **amax_tensor** (Tensor, optional) - Per-group amax tensor.
- **norm_const_tensor** (Tensor, optional) - Normalization constant tensor.
- **prob_tensor** (Tensor) - Per-row gating probabilities tensor.
- **current_stream** (cuda.CUstream) - The CUDA stream to use for execution.
```

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

### Initialize and Execute GroupedGemmGluHadamardSm100

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

Demonstrates how to initialize the GroupedGemmGluHadamardSm100 operation with various sample tensors and parameters, check support, compile, and then execute the operation.

```python
from cudnn import GroupedGemmGluHadamardSm100

op = GroupedGemmGluHadamardSm100(
    sample_a=a,
    sample_b=b,
    sample_c=c,
    sample_d=d,
    sample_sfa=sfa,
    sample_sfb=sfb,
    sample_padded_offsets=padded_offsets,
    sample_alpha=alpha,
    sample_prob=prob,
    sample_amax=amax,
    sample_bias=bias,
    acc_dtype=torch.float32,
    mma_tiler_mn=(256, 256),
    cluster_shape_mn=(2, 1),
    sf_vec_size=16,
    vector_f32=False,
    m_aligned=256,
    act_func="swiglu",
)
assert op.check_support()
op.compile()
op.execute(
    a_tensor=a,
    b_tensor=b,
    c_tensor=c,
    d_tensor=d,
    sfa_tensor=sfa,
    sfb_tensor=sfb,
    padded_offsets=padded_offsets,
    alpha_tensor=alpha,
    prob_tensor=prob,
    amax_tensor=amax,
    bias_tensor=bias,
    current_stream=None,
)
```

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

### Get Autotune Workspace Size

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

Retrieve the workspace size required for autotuning across all available plans. This function is used in conjunction with autotuning to determine the necessary memory allocation.

```cpp
get_autotune_workspace_size() const
```

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

### cuDNN Graph API for Matrix Multiplication

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

Performs the same matrix multiplication as the PyTorch example but using the cuDNN Graph API. This requires explicit definition of data types and tensor names.

```python
import cudnn
import torch

def compare_results(actual: torch.Tensor, expected: torch.Tensor):
    rtol = 1e-2
    atol = 1e-2
    _b, m, n = actual.shape
    assert expected.shape == actual.shape

    # count the number of close elements
    close_mask = torch.isclose(actual, expected, atol=atol, rtol=rtol, equal_nan=True)
    num_el = actual.numel()
    close_cnt = close_mask.detach().sum().cpu().item()
    # find the max diff and location
    max_diff = (actual - expected).abs().max().cpu().item()
    max_diff_idx = (actual - expected).abs().argmax().cpu().item()
    max_diff_idx = (max_diff_idx // (m * n), max_diff_idx % (m * n) // n, max_diff_idx % n)

    print(f"Percentage of close elements: {100 * close_cnt / num_el:.1f}%")
    print(f"Max absolute difference: {max_diff}")
    print(f"At index {list(max_diff_idx)}" 
        f"  GPU={actual[max_diff_idx]}, CPU={expected[max_diff_idx]}")


b, m, n, k = 16, 32, 64, 128
a_dev = torch.randn(b, m, k, device="cuda", dtype=torch.bfloat16)
b_dev = torch.randn(1, k, n, device="cuda", dtype=torch.bfloat16)

# Start of core cuDNN code
with cudnn.Graph(
    io_data_type=torch.bfloat16,
    compute_data_type=torch.float32,
    inputs=["matmul::A", "matmul::B"],
    outputs=["out"],
) as graph:
    c_cudnn = graph.matmul(name="matmul", A=a_dev, B=b_dev)
    c_cudnn.set_name("out").set_output(True)

handle = cudnn.create_handle()
c_dev = graph(a_dev, b_dev, handle=handle)

```

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

### Create Execution Plan with Specific Engine and Knobs

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

Constructs an execution plan using a specified engine ID and a map of knobs. This allows for fine-grained control over plan creation.

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

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

### High-level Wrapper for Grouped GEMM + SwiGLU

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

This snippet demonstrates how to use the `grouped_gemm_swiglu_wrapper_sm100` function for the Grouped GEMM + SwiGLU fusion. It includes example tensor definitions and function call with various parameters.

```APIDOC
## High-level Wrapper

### Description
This function provides a high-level interface to the Grouped GEMM + SwiGLU fusion kernel, simplifying its usage for common scenarios.

### Function Signature
```python
from cudnn import grouped_gemm_swiglu_wrapper_sm100
from cuda.bindings import driver as cuda

outputs = grouped_gemm_swiglu_wrapper_sm100(
    a_tensor, b_tensor, sfa_tensor, sfb_tensor, padded_offsets, alpha_tensor, 
    norm_const_tensor, prob_tensor, acc_dtype, c_dtype, d_dtype, cd_major, 
    mma_tiler_mn, cluster_shape_mn, sf_vec_size, vector_f32, m_aligned, 
    discrete_col_sfd, current_stream
)
```

### Parameters
*   **a_tensor**: Input activation tensor.
*   **b_tensor**: Weight tensor.
*   **sfa_tensor**: Scale factor tensor for A.
*   **sfb_tensor**: Scale factor tensor for B.
*   **padded_offsets**: Cumulative sum of aligned group M sizes.
*   **alpha_tensor**: Per-group scaling factors.
*   **norm_const_tensor**: Normalization constant for FP8 quantization (required when SFD outputs are enabled).
*   **prob_tensor**: Per-row gating probabilities.
*   **acc_dtype**: Accumulator data type (e.g., `torch.float32`).
*   **c_dtype**: Data type for intermediate GEMM output C (e.g., `torch.bfloat16`).
*   **d_dtype**: Data type for the final SwiGLU output D (e.g., `torch.bfloat16`).
*   **cd_major**: Major dimension for C/D tensors ('n' or 'k').
*   **mma_tiler_mn**: MMA tiler shape (M, N).
*   **cluster_shape_mn**: Cluster shape (M, N).
*   **sf_vec_size**: Vector size for scale factors.
*   **vector_f32**: Boolean indicating if vectorization uses f32.
*   **m_aligned**: Aligned M dimension.
*   **discrete_col_sfd**: Boolean indicating if discrete column SFD is used.
*   **current_stream**: CUDA stream object.

### Returns
*   **outputs**: A dictionary or tuple containing the output tensors (C, D, D_row, D_col, SFD_row, SFD_col, amax) as applicable based on input parameters.
```