NVIDIA cuDNN Graph API实战：融合、自动调优与计划复用
Inside NVIDIA’s cuDNN Graph API: Fusion, Autotuning, and Plan Reuse with cuDNN Frontend
提供了一套完整的 cuDNN Graph API 底层开发工作流，包含具体的依赖配置、API调用序列和性能基准测试代码，对需要极致推理优化的工程师极具参考价值。
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.
在本教程中，我们从框架底层逐步讲解 cuDNN Frontend 的图 API：我们将计算描述为操作图，让 cuDNN 选择引擎来运行它，然后由我们自己掌控这一选择。此处构建的每个内核都以相同方式表达：我们按维度和步长声明张量，将操作链接到它们之上，执行五步构建流程（验证、构建操作图、创建执行计划、检查支持情况、构建计划），然后针对指针变体包进行执行。我们在单个 Colab GPU 上运行所有这些操作，并将每个结果与 PyTorch 参考实现进行比较，以便我们既能确认融合的正确性，又能了解其成本。这些主题层层递进，从单个融合卷积扩展到跨引擎配置的自动调优、FP8 风格尾处理、注意力机制、计划序列化、动态形状以及 CUDA 图捕获。
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}")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}")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.
我们首先安装 nvidia-cudnn-frontend，并解决大多数初次运行时会遇到的难题：使 libcudnn.so 对 Frontend 的动态加载器可见。我们强制 PyTorch 优先加载其捆绑的 cuDNN，然后显式预加载共享对象，这样 Frontend 自身的 dlopen 就能解析到进程中已驻留的库。接着，我们报告计算能力，据此选择 bfloat16 或 float16，创建 cuDNN 句柄，并定义张量描述、图构建、工作空间分配以及基于事件的基准测试等辅助函数，供笔记本其余部分复用。
更进一步：量化金融体系
看懂新闻只是起点——沿量化金融路径，把它变成能交付的工程能力