深入 NVIDIA cuDNN Graph API：透過 cuDNN 前端實現融合、自動調優與計畫重用

In this tutorial, we work through the cuDNN Frontend‘s graph API from below the framework: we describe a computation as a graph of operations, let cuDNN pick an engine to run it, and then take control of that choice ourselves.

Every kernel we build here is expressed the same way: we declare tensors by their dimensions and strides, chain operations onto them, run the five-step build pipeline of validate, build operation graph, create execution plans, check support, and build plans, and then execute against a variant pack of pointers.

We run it all on a single Colab GPU, checking each result against a PyTorch reference so we can see both that the fusion is correct and what it costs.

The topics build on each other, moving from a single fused convolution to autotuning across engine configs, FP8-style epilogues, attention, plan serialization, dynamic shapes, and CUDA graph capture.

Copy CodeCopiedUse a different Browserimport 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.printexc(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.isavailable(), "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["CUDNNPATH"] = os.path.dirname(nvidia.cudnn.file) os.environ["LDLIBRARYPATH"] = libdir + ":" + os.environ.

get("LDLIBRARYPATH", "") for so in sorted(glob.glob(os.path.join(libdir, "libcudnn.so"))): try: ctypes.CDLL(so, mode=ctypes.

RTLDGLOBAL) 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.

getdevicecapability() SM = MAJOR 10 + MINOR CUDNNVER = cudnn.backendversion() print(f" GPU : {torch.cuda.getdevicename(0)}") print(f" Compute capability : sm{SM}") print(f" Torch / CUDA : {torch.version} / {torch.version.

cuda}") print(f" cuDNN backend : {CUDNNVER}") try: print(f" cuDNN version str : {cudnn.backendversionstring()}") except Exception: pass DTYPE = torch.bfloat16 if SM >= 80 else torch.float16 HASSDPA = SM >= 80 print(f" Working dtype : {DTYPE}") print(f" Fused SDPA usable : {HASSDPA}") HANDLE = cudnn.

createhandle() TORCH2CUDNN = { torch.float16: cudnn.datatype.HALF, torch.bfloat16: cudnn.datatype.BFLOAT16, torch.float32: cudnn.datatype.FLOAT, torch.int32: cudnn.datatype.INT32, torch.int64: cudnn.datatype.INT64, torch.int8: cudnn.datatype.INT8, torch.uint8: cudnn.datatype.

UINT8, } def tensorof(graph, t, name): return graph.tensor( name=name, dim=list(t.size()), stride=list(t.stride()), datatype=TORCH2CUDNN[t.dtype], ) def scalarof(graph, name): return graph.tensor( name=name, dim=[1, 1, 1], stride=[1, 1, 1], datatype=cudnn.datatype.

FLOAT, ispassbyvalue=True, ) def build(graph, heur=None, policy=None): heur = heur or [cudnn.heurmode.A, cudnn.heurmode.FALLBACK] graph.validate() graph.buildoperationgraph() graph.createexecutionplans(heur) graph.checksupport() if policy is None: graph.buildplans() else: graph.

buildplans(policy) return graph def workspacefor(graph): n = graph.getworkspacesize() 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.elapsedtime(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}") We start by installing nvidia-cudnn-frontend and solving the problem that trips up most first runs: making libcudnn.so visible to the frontend’s dynamic loader.

We force PyTorch to load its bundled cuDNN first and then preload the shared objects explicitly, so the frontend’s own dlopen resolves against a library already resident in the process.

We then report the compute capability, pick bfloat16 or float16 accordingly, create the cuDNN handle, and define the helpers for tensor description, graph building, workspace allocation, and event-based benchmarking that the rest of the notebook reuses.

Copy CodeCopiedUse a different BrowserN, 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 CONVFLOPS = 2 N K P Q C R S CONVSTATE = {} @section("2.

Fused Conv -> Bias -> ReLU") def convfusion(): x = torch.randn(N, C, H, W, device=DEV, dtype=DTYPE).to(memoryformat=torch.channelslast) w = torch.randn(K, C, R, S, device=DEV, dtype=DTYPE).to(memoryformat=torch.channelslast) b = torch.randn(1, K, 1, 1, device=DEV, dtype=DTYPE) y = torch.

empty(N, K, P, Q, device=DEV, dtype=DTYPE).to(memoryformat=torch.channelslast) g = cudnn.pygraph( handle=HANDLE, name="convbiasrelu", iodatatype=TORCH2CUDNN[DTYPE], intermediatedatatype=cudnn.datatype.FLOAT, computedatatype=cudnn.datatype.

FLOAT, ) X = tensorof(g, x, "X") Wt = tensorof(g, w, "W") Bt = tensorof(g, b, "bias") conv = g.convfprop( image=X, weight=Wt, padding=[PAD, PAD], stride=[STR, STR], dilation=[DIL, DIL], computedatatype=cudnn.datatype.FLOAT, ) biased = g.bias(input=conv, bias=Bt) Y = g.relu(input=biased) Y.

setoutput(True).setdatatype(TORCH2CUDNN[DTYPE]) Y.setdim(list(y.size())).setstride(list(y.stride())) t0 = time.perfcounter() build(g) buildms = (time.perfcounter() - t0) 1e3 ws = workspacefor(g) pack = {X: x, Wt: w, Bt: b, Y: y} g.execute(pack, ws) torch.cuda.synchronize() ref = torch.relu(torch.nn.

functional.conv2d(x, w, bias=b.flatten(), padding=PAD)) err = (y.float() - ref.float()).abs().max().item() scale = ref.float().abs().max().item() print(f" problem : N{N} C{C} {H}x{W} -> K{K} {R}x{S} ({DTYPE})") print(f" build : {buildms:.1f} ms workspace: {ws.numel()/1024:.

1f} KiB") print(f" max |err|: {err:.4f} (ref max {scale:.2f}, rel {err/max(scale,1e-9):.2e})") assert err / max(scale, 1e-9) < 5e-2, "numerical mismatch vs PyTorch" mscudnn = bench(lambda: g.execute(pack, ws)) mstorch = bench(lambda: torch.relu( torch.nn.functional.conv2d(x, w, bias=b.

flatten(), padding=PAD))) print() report("cuDNN FE (single fused kernel)", mscudnn, CONVFLOPS) report("PyTorch (conv+bias, then relu)", mstorch, CONVFLOPS) print(f" speedup: {mstorch/mscudnn:.2f}x") CONVSTATE.update(graph=g, pack=pack, ws=ws, x=x, w=w, b=b, y=y) return f"{mscudnn:.

3f} ms, {tflops(CONVFLOPS, mscudnn):.1f} TFLOP/s" convfusion() We build our first graph, a convolution followed by a bias add and a ReLU, all fused into a single kernel.

We keep every tensor in channels_last because that is what gives cuDNN the NHWC strides its tensor-core engines want, and we pin the output dimensions and strides explicitly so the result is written back in the same layout.We validate the output against torch.nn.functional.

conv2d, then benchmark the fused graph against PyTorch running the convolution and activation as separate kernels.Copy CodeCopiedUse a diffe