Quick Start Guide to cuDNN

Start from a small PyTorch linear layer on a GPU:

import torch
import torch.nn as nn
b, m, n, k = 16, 32, 64, 128
x = torch.randn(b, m, k, device="cuda", dtype=torch.bfloat16)
linear = nn.Linear(k, n, bias=False, device="cuda", dtype=torch.bfloat16)
y = linear(x)
# same:
# y = torch.matmul(x, linear.weight.T)

The tensors have shape [B, M, K] and a projection along the last axis to width N—a pattern common in LLMs (batch B, sequence M, hidden K, output width N). Conceptually, the operation is a batched matrix multiplication. PyTorch expresses that concept compactly and dispatches to its own kernels.

The PyTorch block shown previously is a baseline for the same math, not a suggestion that cuDNN should match it line-for-line. Frameworks typically hide the graph construction and much of the dtype and the layout detail that the cuDNN frontend asks you to state explicitly.

The following listing runs the same matrix multiplication through the cuDNN graph API.

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)
# End of core cuDNN code
c_ref = torch.matmul(a_dev.to(torch.float32), b_dev.to(torch.float32)).to(torch.bfloat16)
compare_results(c_dev, c_ref)

The cuDNN example stays at the matrix multiplication level (no nn.Module wrapper). It is more explicit than PyTorch about tensor roles and dtypes.

Core cuDNN Path

The with cudnn.Graph(...) context through c_dev = graph(...).

Supporting Checks

compare_results, the float32 torch.matmul reference, and the printout exist only to show that c_dev matches PyTorch within a loose numerical tolerance. They exist only as a confidence check, not as production overhead you must carry everywhere.

Beyond a single matrix multiplication, the graph-based API is how the frontend composes multiple operations and reuses a built graph. That workflow does not map one-to-one to a one-line nn.Linear call.

Graph Basics

A cuDNN graph is a DAG of ops over tensors. A built graph maps to CUDA work on the device. Support depends on the GPU and the library version. When defining a graph, specify at least:

• Which tensors are the graph outputs
• The I/O dtypes for the graph boundaries
• The per-node compute dtype (often FP32 accumulation for matrix multiplication)
• The intermediate dtypes between nodes (only when the graph has multiple nodes)

The sample sets default I/O and compute dtypes on the graph context. With a single matrix multiplication node, no intermediate dtype is needed. Shapes and strides for a_dev and b_dev are inferred from the PyTorch tensors passed into the graph builder. At definition time, only metadata is used. The tensor values are read at execution.

Build and Validate

When the cudnn.Graph context exits, the frontend validates and builds the graph. Invalid DAGs or unsupported patterns raise at exit. Output metadata for c_cudnn is filled in during that step even if it was not set explicitly.

Execution and Reuse

After the build, the graph object can be called repeatedly like a function. The inputs and outputs lists fix argument order. String names (for example matmul::A for argument A of node matmul, and out for the marked output) wire the graph ports to the call arguments. PyTorch tensors used only at build time still supply live buffers at run time.

The graph setup runs on the host. Reuse the built graph when shapes, layouts, and dtypes stay fixed, so setup cost is amortized. Pass a cudnn handle into the call (for example from cudnn.create_handle()) to select the device and stream context, or pass handle="auto" for a graph-owned handle (less sharing across graphs).

compare_results() reports how closely c_dev matches c_ref. c_ref uses float32 accumulation then casts back to bfloat16 so the comparison aligns with the graph’s mixed-precision matrix multiplication.