AI 服務暫時不可用，以下為來源正文，待恢復後補全翻譯。
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 Browser 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: 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 = g.conv_fprop( image=X, weight=Wt, padding=[PAD, PAD], stride=[STR, STR], dilation=[DIL, DIL], compute_data_type=cudnn.data_type.FLOAT, ) biased = g.bias(input=conv, bias=Bt) Y = g.relu(input=biased) Y.set_output(True).set_data_type(TORCH2CUDNN[DTYPE]) Y.set_dim(list(y.size())).set_stride(list(y.stride())) t0 = time.perf_counter() build(g) build_ms = (time.perf_counter() - t0) * 1e3 ws = workspace_for(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 : {build_ms:.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) 3d}: {ms:8.3f} ms " f"{tflops(CONV_FLOPS, ms):7.2f} TFLOP/s ws={ws_sz/1024:8.1f} KiB") except Exception as e: print(f" plan {i:>3d}: unusable ({type(e).name})") assert timings, "no plan executed" timings.sort() best_ms, best_i, best_ws = timings[0] worst_ms = timings[-1][0] print(f"\n fastest = plan {best_i} @ {best_ms:.3f} ms") print(f" slowest = {worst_ms:.3f} ms -> {worst_ms/best_ms:.1f}x spread across engines") print(" Takeaway: heuristics are good, but for a hot shape you ship the") print(" autotuned index (or the serialized plan from section 6).") return f"best plan {best_i} @ {best_ms:.3f} ms ({worst_ms/best_ms:.1f}x spread)" autotune() We rebuild the same convolution but stop trusting the heuristic, asking for plans from heuristic modes A, B, and FALLBACK and compiling all of them with build_plan_policy.ALL. We then walk the plan list, build each config, allocate its specific workspace, and time it with execute_plan_at_index, printing throughput and workspace size for every candidate. The spread between the fastest and slowest engine is the point of the exercise, because it tells us how much we gain by shipping an autotuned index instead of accepting the default pick. Copy CodeCopiedUse a different Browser @section("4. Matmul -> scale -> bias -> activation -> AMAX") def matmul_epilogue(): Bsz, M, Kd, Nd = 16, 512, 1024, 512 MM_FLOPS = 2 * Bsz * M * Nd * Kd a = torch.randn(Bsz, M, Kd, device=DEV, dtype=DTYPE) bm = torch.randn(Bsz, Kd, Nd, device=DEV, dtype=DTYPE) bias = torch.randn(1, 1, Nd, device=DEV, dtype=DTYPE) out = torch.empty(Bsz, M, Nd, device=DEV, dtype=DTYPE) amax = torch.empty(1, 1, 1, device=DEV, dtype=torch.float32) alpha_val = 0.125 alpha = torch.full((1, 1, 1), alpha_val, dtype=torch.float32) g = cudnn.pygraph( handle=HANDLE, name="matmul_epilogue", io_data_type=TORCH2CUDNN[DTYPE], intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT, ) A = tensor_of(g, a, "A") Bt = tensor_of(g, bm, "B") BIAS = tensor_of(g, bias, "bias") ALPHA = scalar_of(g, "alpha") acc = g.matmul(A=A, B=Bt, compute_data_type=cudnn.data_type.FLOAT) scaled = g.mul(a=acc, b=ALPHA) biased = g.bias(input=scaled, bias=BIAS) act_name = "relu" if hasattr(g, "gelu"): try: act = g.gelu(input=biased) act_name = "gelu" except Exception: act = g.relu(input=biased) else: act = g.relu(input=biased) print(f" activation used: {act_name}") OUT = act OUT.set_output(True).set_data_type(TORCH2CUDNN[DTYPE]) have_amax = True try: AMAX = g.reduction(input=act, mode=cudnn.reduction_mode.AMAX, compute_data_type=cudnn.data_type.FLOAT) AMAX.set_output(True).set_data_type(cudnn.data_type.FLOAT) AMAX.set_dim([1, 1, 1]).set_stride([1, 1, 1]) except Exception as e: have_amax = False print(f" (AMAX reduction unavailable here: {e})") build(g) ws = workspace_for(g) pack = {A: a, Bt: bm, BIAS: bias, ALPHA: alpha, OUT: out} if have_amax: pack[AMAX] = amax g.execute(pack, ws) torch.cuda.synchronize() ref = torch.matmul(a.float(), bm.float()) * alpha_val + bias.float() ref = torch.nn.functional.gelu(ref) if act_name == "gelu" else torch.relu(ref) rel = ((out.float() - ref).abs().max() / ref.abs().max()).item() print(f" shape : ({Bsz},{M},{Kd}) x ({Bsz},{Kd},{Nd})") print(f" rel err : {rel:.2e}") if have_amax: print(f" fused AMAX {amax.item():.4f} vs torch {ref.abs().max().item():.4f}") ms = bench(lambda: g.execute(pack, ws)) def torch_ref(): r = torch.baddbmm(bias.expand(Bsz, M, Nd), a, bm, beta=1.0, alpha=alpha_val) r = torch.nn.functional.gelu(r) if act_name == "gelu" else torch.relu(r) return r.abs().amax() ms_t = bench(torch_ref) print() report("cuDNN FE (one fused kernel)", ms, MM_FLOPS) report("PyTorch (bmm + act + amax)", ms_t, MM_FLOPS) print(f" speedup: {ms_t/ms:.2f}x -- the win is the epilogue traffic, not the GEMM") return f"{ms:.3f} ms, {tflops(MM_FLOPS, ms):.1f} TFLOP/s, {ms_t/ms:.2f}x vs torch" matmul_epilogue() We move to a batched matmul and hang a full epilogue off it: an alpha scale supplied as a pass-by-value host scalar, a bias add, an activation, and an AMAX reduction over the result. The AMAX in the same kernel is the pattern that FP8 training relies on, since it collects the scale factor for the next quantization step without a second pass over the output. We compare against a PyTorch chain of baddbmm, activation, and amax, which makes clear that the speedup comes from eliminating epilogue memory traffic rather than from a faster GEMM. Copy CodeCopiedUse a different Browser @section("5. SDPA (Flash Attention) with causal masking") def sdpa_demo(): if not HAS_SDPA: raise RuntimeError(f"fused SDPA needs SM80+ (Ampere), this GPU is sm_{SM}") b, h, s, d = 4, 16, 1024, 64 scale = 1.0 / math.sqrt(d) SDPA_FLOPS = 4 * b * h * s * s * d * 0.5 q = torch.randn(b, h, s, d, device=DEV, dtype=DTYPE) k = torch.randn(b, h, s, d, device=DEV, dtype=DTYPE) v = torch.randn(b, h, s, d, device=DEV, dtype=DTYPE) o = torch.empty(b, h, s, d, device=DEV, dtype=DTYPE) g = cudnn.pygraph( handle=HANDLE, name="sdpa", io_data_type=TORCH2CUDNN[DTYPE], intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT, ) Q, Kt, V = tensor_of(g, q, "Q"), tensor_of(g, k, "K"), tensor_of(g, v, "V") causal = True try: O, _stats = g.sdpa(name="sdpa", q=Q, k=Kt, v=V, is_inference=True, attn_scale=scale, use_causal_mask=True) except TypeError: try: O, _stats = g.sdpa(name="sdpa", q=Q, k=Kt, v=V, is_inference=True, attn_scale=scale, diagonal_alignment=cudnn.diagonal_alignment.TOP_LEFT, right_bound=0) except Exception: causal = False O, _stats = g.sdpa(name="sdpa", q=Q, k=Kt, v=V, is_inference=True, attn_scale=scale) print(f" causal masking: {causal}") O.set_output(True).set_data_type(TORCH2CUDNN[DTYPE]) O.set_dim(list(o.size())).set_stride(list(o.stride())) build(g) ws = workspace_for(g) pack = {Q: q, Kt: k, V: v, O: o} g.execute(pack, ws) torch.cuda.synchronize() ref = torch.nn.functional.scaled_dot_product_attention(q, [truncated for AI cost control]