Developers can now construct deep learning graphs directly in Python, letting cuDNN select the execution engine while retaining control over that choice.
The workflow
The approach treats computation as a graph of operations. Users declare tensors by their dimensions and strides, chain operations onto them, and run a five-step pipeline: validate, build the operation graph, create execution plans, check support, and build plans. Execution happens against a variant pack of pointers. The tutorial runs this on a single Colab GPU, comparing results against PyTorch references to verify correctness and measure cost.
Topics progress from a single fused convolution to autotuning across engine configs, FP8-style epilogues, attention mechanisms, plan serialization, dynamic shapes, and CUDA graph capture.
Setup and environment
The first step installs nvidia-cudnn-frontend. This solves the common issue of the dynamic loader failing to find libcudnn.so. The code forces PyTorch to load its bundled cuDNN first, then preloads the shared objects explicitly so the frontend resolves against a library already resident in the process.
The script reports the compute capability, selects bfloat16 or float16 based on hardware, creates a cuDNN handle, and defines helpers for tensor description, graph building, workspace allocation, and event-based benchmarking.
import os
import sys
import glob
import math
import time
import ctypes
import traceback
import subprocess
RESULTS = {}
def banner(title):
print("\n" + "=" * 78)
print(title)
print("=" * 78)
def section(name):
def wrap(fn):
def run(*a, **kw):
banner(name)
try:
out = fn(*a, **kw)
RESULTS[name] = out if isinstance(out, str) else "ok"
return out
except Exception as e:
RESULTS[name] = f"SKIPPED / FAILED -> {type(e).__name__}: {e}"
print(f"\n[!] {name} did not complete: {type(e).__name__}: {e}")
traceback.print_exc(limit=3)
return None
return run
return wrap
banner("0. Install nvidia-cudnn-frontend and locate libcudnn")
subprocess.run(
[sys.executable, "-m", "pip", "install", "-q", "nvidia-cudnn-frontend"],
check=True,
)
import torch
assert torch.cuda.is_available(), "No GPU. Runtime -> Change runtime type -> GPU."
torch.backends.cudnn.enabled = True
_ = torch.nn.functional.conv2d(
torch.randn(1, 1, 8, 8, device="cuda"), torch.randn(1, 1, 3, 3, device="cuda")
)
torch.cuda.synchronize()
try:
import nvidia.cudnn
_libdir = os.path.join(os.path.dirname(nvidia.cudnn.__file__), "lib")
os.environ["CUDNN_PATH"] = os.path.dirname(nvidia.cudnn.__file__)
os.environ["LD_LIBRARY_PATH"] = _libdir + ":" + os.environ.get("LD_LIBRARY_PATH", "")
for _so in sorted(glob.glob(os.path.join(_libdir, "libcudnn*.so*"))):
try:
ctypes.CDLL(_so, mode=ctypes.RTLD_GLOBAL)
except OSError:
pass
except Exception as _e:
print(f" (no pip cuDNN package found, relying on system cuDNN: {_e})")
import cudnn
print(" cuDNN frontend imported successfully.")
banner("1. Environment")
DEV = torch.device("cuda")
MAJOR, MINOR = torch.cuda.get_device_capability()
SM = MAJOR * 10 + MINOR
CUDNN_VER = cudnn.backend_version()
print(f" GPU : {torch.cuda.get_device_name(0)}")
print(f" Compute capability : sm_{SM}")
print(f" Torch / CUDA : {torch.__version__} / {torch.version.cuda}")
print(f" cuDNN backend : {CUDNN_VER}")
try:
print(f" cuDNN version str : {cudnn.backend_version_string()}")
except Exception:
pass
DTYPE = torch.bfloat16 if SM >= 80 else torch.float16
HAS_SDPA = SM >= 80
print(f" Working dtype : {DTYPE}")
print(f" Fused SDPA usable : {HAS_SDPA}")
HANDLE = cudnn.create_handle()
TORCH2CUDNN = {
torch.float16: cudnn.data_type.HALF,
torch.bfloat16: cudnn.data_type.BFLOAT16,
torch.float32: cudnn.data_type.FLOAT,
torch.int32: cudnn.data_type.INT32,
torch.int64: cudnn.data_type.INT64,
torch.int8: cudnn.data_type.INT8,
torch.uint8: cudnn.data_type.UINT8,
}
def tensor_of(graph, t, name):
return graph.tensor(
name=name,
dim=list(t.size()),
stride=list(t.stride()),
data_type=TORCH2CUDNN[t.dtype],
)
def scalar_of(graph, name):
return graph.tensor(
name=name,
dim=[1, 1, 1],
stride=[1, 1, 1],
data_type=cudnn.data_type.FLOAT,
is_pass_by_value=True,
)
def build(graph, heur=None, policy=None):
heur = heur or [cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]
graph.validate()
graph.build_operation_graph()
graph.create_execution_plans(heur)
graph.check_support()
if policy is None:
graph.build_plans()
else:
graph.build_plans(policy)
return graph
def workspace_for(graph):
n = graph.get_workspace_size()
return torch.empty(max(n, 1), device=DEV, dtype=torch.uint8)
def bench(fn, warmup=10, iters=50):
for _ in range(warmup):
fn()
torch.cuda.synchronize()
s, e = torch.cuda.Event(True), torch.cuda.Event(True)
s.record()
for _ in range(iters):
fn()
e.record()
torch.cuda.synchronize()
return s.elapsed_time(e) / iters
def tflops(flops, ms):
return flops / (ms * 1e-3) / 1e12
def report(tag, ms, flops=None):
extra = f" ({tflops(flops, ms):7.2f} TFLOP/s)" if flops else ""
print(f" {tag:<34s} {ms:8.3f} ms{extra}")
Fused convolution
The code builds a graph for a convolution followed by a bias add and a ReLU, all fused into a single kernel. Every tensor uses channels_last memory format. This provides the NHWC strides required by cuDNN tensor-core engines. The output dimensions and strides are pinned explicitly so the result writes back in the same layout.
Validation against torch.nn.functional.conv2d confirms numerical accuracy. The benchmark compares the fused graph against PyTorch running the convolution and activation as separate kernels.
N, C, H, W = 32, 128, 56, 56
K, R, S = 256, 3, 3
PAD, STR, DIL = 1, 1, 1
P = (H + 2 * PAD - DIL * (R - 1) - 1) // STR + 1
Q = (W + 2 * PAD - DIL * (S - 1) - 1) // STR + 1
CONV_FLOPS = 2 * N * K * P * Q * C * R * S
CONV_STATE = {}
@section("2. Fused Conv -> Bias -> ReLU")
def conv_fusion():
x = torch.randn(N, C, H, W, device=DEV, dtype=DTYPE).to(memory_format=torch.channels_last)
w = torch.randn(K, C, R, S, device=DEV, dtype=DTYPE).to(memory_format=torch.channels_last)
b = torch.randn(1, K, 1, 1, device=DEV, dtype=DTYPE)
y = torch.empty(N, K, P, Q, device=DEV, dtype=DTYPE).to(memory_format=torch.channels_last)
g = cudnn.pygraph(
handle=HANDLE,
name="conv_bias_relu",
io_data_type=TORCH2CUDNN[DTYPE],
intermediate_data_type=cudnn.data_type.FLOAT,
compute_data_type=cudnn.data_type.FLOAT,
)
X = tensor_of(g, x, "X")
Wt = tensor_of(g, w, "W")
Bt = tensor_of(g, b, "bias")
conv = gSource Read original →Related articles
