Inside NVIDIA’s cuDNN Graph API: Fusion, Autotuning, and Plan Reuse with cuDNN Frontend
In this tutorial, we work by the cuDNN Frontend‘s graph API from under the framework: we describe a computation as a graph of operations, let cuDNN choose an engine to run it, and then take management of that selection ourselves. Every kernel we construct right here is expressed the identical manner: we declare tensors by their dimensions and strides, chain operations onto them, run the five-step construct pipeline of validate, construct operation graph, create execution plans, verify help, and construct plans, and then execute in opposition to a variant pack of pointers. We run all of it on a single Colab GPU, checking every consequence in opposition to a PyTorch reference so we will see each that the fusion is appropriate and what it prices. The matters construct on one another, transferring from a single fused convolution to autotuning throughout engine configs, FP8-style epilogues, consideration, plan serialization, dynamic shapes, and CUDA graph seize.
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 part(identify):
def wrap(fn):
def run(*a, **kw):
banner(identify)
strive:
out = fn(*a, **kw)
RESULTS[name] = out if isinstance(out, str) else "okay"
return out
besides Exception as e:
RESULTS[name] = f"SKIPPED / FAILED -> {kind(e).__name__}: {e}"
print(f"n[!] {identify} didn't full: {kind(e).__name__}: {e}")
traceback.print_exc(restrict=3)
return None
return run
return wrap
banner("0. Install nvidia-cudnn-frontend and find libcudnn")
subprocess.run(
[sys.executable, "-m", "pip", "install", "-q", "nvidia-cudnn-frontend"],
verify=True,
)
import torch
assert torch.cuda.is_available(), "No GPU. Runtime -> Change runtime kind -> GPU."
torch.backends.cudnn.enabled = True
_ = torch.nn.practical.conv2d(
torch.randn(1, 1, 8, 8, system="cuda"), torch.randn(1, 1, 3, 3, system="cuda")
)
torch.cuda.synchronize()
strive:
import nvidia.cudnn
_libdir = os.path.be part of(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.be part of(_libdir, "libcudnn*.so*"))):
strive:
ctypes.CDLL(_so, mode=ctypes.RTLD_GLOBAL)
besides OSError:
move
besides Exception as _e:
print(f" (no pip cuDNN bundle discovered, counting on system cuDNN: {_e})")
import cudnn
print(" cuDNN frontend imported efficiently.")
banner("1. Environment")
DEV = torch.system("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 functionality : sm_{SM}")
print(f" Torch / CUDA : {torch.__version__} / {torch.model.cuda}")
print(f" cuDNN backend : {CUDNN_VER}")
strive:
print(f" cuDNN model str : {cudnn.backend_version_string()}")
besides Exception:
move
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, identify):
return graph.tensor(
identify=identify,
dim=listing(t.dimension()),
stride=listing(t.stride()),
data_type=TORCH2CUDNN[t.dtype],
)
def scalar_of(graph, identify):
return graph.tensor(
identify=identify,
dim=[1, 1, 1],
stride=[1, 1, 1],
data_type=cudnn.data_type.FLOAT,
is_pass_by_value=True,
)
def construct(graph, heur=None, coverage=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 coverage is None:
graph.build_plans()
else:
graph.build_plans(coverage)
return graph
def workspace_for(graph):
n = graph.get_workspace_size()
return torch.empty(max(n, 1), system=DEV, dtype=torch.uint8)
def bench(fn, warmup=10, iters=50):
for _ in vary(warmup):
fn()
torch.cuda.synchronize()
s, e = torch.cuda.Event(True), torch.cuda.Event(True)
s.document()
for _ in vary(iters):
fn()
e.document()
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):
additional = f" ({tflops(flops, ms):7.2f} TFLOP/s)" if flops else ""
print(f" {tag:<34s} {ms:8.3f} ms{additional}")
We begin by putting in nvidia-cudnn-frontend and fixing the issue that journeys up most first runs: making libcudnn.so seen to the frontend’s dynamic loader. We power PyTorch to load its bundled cuDNN first and then preload the shared objects explicitly, so the frontend’s personal dlopen resolves in opposition to a library already resident within the course of. We then report the compute functionality, choose bfloat16 or float16 accordingly, create the cuDNN deal with, and outline the helpers for tensor description, graph constructing, workspace allocation, and event-based benchmarking that the remainder of the pocket book reuses.
N, C, H, W = 32, 128, 56, 56
Ok, 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 * Ok * P * Q * C * R * S
CONV_STATE = {}
@part("2. Fused Conv -> Bias -> ReLU")
def conv_fusion():
x = torch.randn(N, C, H, W, system=DEV, dtype=DTYPE).to(memory_format=torch.channels_last)
w = torch.randn(Ok, C, R, S, system=DEV, dtype=DTYPE).to(memory_format=torch.channels_last)
b = torch.randn(1, Ok, 1, 1, system=DEV, dtype=DTYPE)
y = torch.empty(N, Ok, P, Q, system=DEV, dtype=DTYPE).to(memory_format=torch.channels_last)
g = cudnn.pygraph(
deal with=HANDLE,
identify="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(
picture=X, weight=Wt,
padding=[PAD, PAD], stride=[STR, STR], dilation=[DIL, DIL],
compute_data_type=cudnn.data_type.FLOAT,
)
biased = g.bias(enter=conv, bias=Bt)
Y = g.relu(enter=biased)
Y.set_output(True).set_data_type(TORCH2CUDNN[DTYPE])
Y.set_dim(listing(y.dimension())).set_stride(listing(y.stride()))
t0 = time.perf_counter()
construct(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.practical.conv2d(x, w, bias=b.flatten(), padding=PAD))
err = (y.float() - ref.float()).abs().max().merchandise()
scale = ref.float().abs().max().merchandise()
print(f" downside : N{N} C{C} {H}x{W} -> Ok{Ok} {R}x{S} ({DTYPE})")
print(f" construct : {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) < 5e-2, "numerical mismatch vs PyTorch"
ms_cudnn = bench(lambda: g.execute(pack, ws))
ms_torch = bench(lambda: torch.relu(
torch.nn.practical.conv2d(x, w, bias=b.flatten(), padding=PAD)))
print()
report("cuDNN FE (single fused kernel)", ms_cudnn, CONV_FLOPS)
report("PyTorch (conv+bias, then relu)", ms_torch, CONV_FLOPS)
print(f" speedup: {ms_torch/ms_cudnn:.2f}x")
CONV_STATE.replace(graph=g, pack=pack, ws=ws, x=x, w=w, b=b, y=y)
return f"{ms_cudnn:.3f} ms, {tflops(CONV_FLOPS, ms_cudnn):.1f} TFLOP/s"
conv_fusion()
We construct our first graph, a convolution adopted by a bias add and a ReLU, all fused right into a single kernel. We hold each tensor in channels_last as a result of that’s what offers cuDNN the NHWC strides its tensor-core engines need, and we pin the output dimensions and strides explicitly so the result’s written again in the identical structure. We validate the output in opposition to torch.nn.practical.conv2d, then benchmark the fused graph in opposition to PyTorch working the convolution and activation as separate kernels.
@part("3. Autotuning: construct ALL plans, time every engine config")
def autotune():
x, w, b, y = CONV_STATE["x"], CONV_STATE["w"], CONV_STATE["b"], CONV_STATE["y"]
g = cudnn.pygraph(
deal with=HANDLE, identify="conv_autotune",
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")
Y = g.relu(enter=g.bias(
enter=g.conv_fprop(picture=X, weight=Wt, padding=[PAD, PAD],
stride=[STR, STR], dilation=[DIL, DIL],
compute_data_type=cudnn.data_type.FLOAT),
bias=Bt))
Y.set_output(True).set_data_type(TORCH2CUDNN[DTYPE])
Y.set_dim(listing(y.dimension())).set_stride(listing(y.stride()))
g.validate()
g.build_operation_graph()
g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.B, cudnn.heur_mode.FALLBACK])
g.check_support()
g.build_plans(cudnn.build_plan_policy.ALL)
n_plans = g.get_execution_plan_count()
print(f" {n_plans} candidate engine configs survived help checksn")
pack = {X: x, Wt: w, Bt: b, Y: y}
timings = []
for i in vary(n_plans):
strive:
g.build_plan_at_index(i)
ws_sz = max(g.get_workspace_size_plan_at_index(i), 1)
ws = torch.empty(ws_sz, system=DEV, dtype=torch.uint8)
ms = bench(lambda: g.execute_plan_at_index(pack, ws, i), warmup=3, iters=15)
timings.append((ms, i, ws_sz))
print(f" plan {i:>3d}: {ms:8.3f} ms "
f"{tflops(CONV_FLOPS, ms):7.2f} TFLOP/s ws={ws_sz/1024:8.1f} KiB")
besides Exception as e:
print(f" plan {i:>3d}: unusable ({kind(e).__name__})")
assert timings, "no plan executed"
timings.type()
best_ms, best_i, best_ws = timings[0]
worst_ms = timings[-1][0]
print(f"n quickest = plan {best_i} @ {best_ms:.3f} ms")
print(f" slowest = {worst_ms:.3f} ms -> {worst_ms/best_ms:.1f}x unfold throughout engines")
print(" Takeaway: heuristics are good, however for a sizzling form you ship the")
print(" autotuned index (or the serialized plan from part 6).")
return f"greatest plan {best_i} @ {best_ms:.3f} ms ({worst_ms/best_ms:.1f}x unfold)"
autotune()
We rebuild the identical convolution however cease 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 stroll the plan listing, construct every config, allocate its particular workspace, and time it with execute_plan_at_index, printing throughput and workspace dimension for each candidate. The unfold between the quickest and slowest engine is the purpose of the train, as a result of it tells us how a lot we achieve by delivery an autotuned index as an alternative of accepting the default choose.
@part("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, system=DEV, dtype=DTYPE)
bm = torch.randn(Bsz, Kd, Nd, system=DEV, dtype=DTYPE)
bias = torch.randn(1, 1, Nd, system=DEV, dtype=DTYPE)
out = torch.empty(Bsz, M, Nd, system=DEV, dtype=DTYPE)
amax = torch.empty(1, 1, 1, system=DEV, dtype=torch.float32)
alpha_val = 0.125
alpha = torch.full((1, 1, 1), alpha_val, dtype=torch.float32)
g = cudnn.pygraph(
deal with=HANDLE, identify="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(enter=scaled, bias=BIAS)
act_name = "relu"
if hasattr(g, "gelu"):
strive:
act = g.gelu(enter=biased)
act_name = "gelu"
besides Exception:
act = g.relu(enter=biased)
else:
act = g.relu(enter=biased)
print(f" activation used: {act_name}")
OUT = act
OUT.set_output(True).set_data_type(TORCH2CUDNN[DTYPE])
have_amax = True
strive:
AMAX = g.discount(enter=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])
besides Exception as e:
have_amax = False
print(f" (AMAX discount unavailable right here: {e})")
construct(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.practical.gelu(ref) if act_name == "gelu" else torch.relu(ref)
rel = ((out.float() - ref).abs().max() / ref.abs().max()).merchandise()
print(f" form : ({Bsz},{M},{Kd}) x ({Bsz},{Kd},{Nd})")
print(f" rel err : {rel:.2e}")
if have_amax:
print(f" fused AMAX {amax.merchandise():.4f} vs torch {ref.abs().max().merchandise():.4f}")
ms = bench(lambda: g.execute(pack, ws))
def torch_ref():
r = torch.baddbmm(bias.increase(Bsz, M, Nd), a, bm, beta=1.0, alpha=alpha_val)
r = torch.nn.practical.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 site visitors, 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 transfer to a batched matmul and hold a full epilogue off it: an alpha scale equipped as a pass-by-value host scalar, a bias add, an activation, and an AMAX discount over the consequence. The AMAX in the identical kernel is the sample that FP8 coaching depends on, because it collects the dimensions issue for the subsequent quantization step and not using a second move over the output. We evaluate in opposition to a PyTorch chain of baddbmm, activation, and amax, which makes clear that the speedup comes from eliminating epilogue reminiscence site visitors reasonably than from a quicker GEMM.
@part("5. SDPA (Flash Attention) with causal masking")
def sdpa_demo():
if not HAS_SDPA:
elevate RuntimeError(f"fused SDPA wants 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, system=DEV, dtype=DTYPE)
ok = torch.randn(b, h, s, d, system=DEV, dtype=DTYPE)
v = torch.randn(b, h, s, d, system=DEV, dtype=DTYPE)
o = torch.empty(b, h, s, d, system=DEV, dtype=DTYPE)
g = cudnn.pygraph(
deal with=HANDLE, identify="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, ok, "Ok"), tensor_of(g, v, "V")
causal = True
strive:
O, _stats = g.sdpa(identify="sdpa", q=Q, ok=Kt, v=V,
is_inference=True, attn_scale=scale, use_causal_mask=True)
besides TypeError:
strive:
O, _stats = g.sdpa(identify="sdpa", q=Q, ok=Kt, v=V,
is_inference=True, attn_scale=scale,
diagonal_alignment=cudnn.diagonal_alignment.TOP_LEFT,
right_bound=0)
besides Exception:
causal = False
O, _stats = g.sdpa(identify="sdpa", q=Q, ok=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(listing(o.dimension())).set_stride(listing(o.stride()))
construct(g)
ws = workspace_for(g)
pack = {Q: q, Kt: ok, V: v, O: o}
g.execute(pack, ws)
torch.cuda.synchronize()
ref = torch.nn.practical.scaled_dot_product_attention(q, ok, v, is_causal=causal, scale=scale)
rel = ((o.float() - ref.float()).abs().max() / ref.float().abs().max()).merchandise()
print(f" form : b{b} h{h} s{s} d{d} workspace {ws.numel()/1024:.1f} KiB")
print(f" rel err : {rel:.2e}")
ms = bench(lambda: g.execute(pack, ws))
ms_t = bench(lambda: torch.nn.practical.scaled_dot_product_attention(
q, ok, v, is_causal=causal, scale=scale))
print()
report("cuDNN FE SDPA", ms, SDPA_FLOPS)
report("torch SDPA (backend's selection)", ms_t, SDPA_FLOPS)
print(" Note: torch might already be dispatching to cuDNN or FlashAttention,")
print(" so parity right here is the anticipated, wholesome end result.")
return f"{ms:.3f} ms, {tflops(SDPA_FLOPS, ms):.1f} TFLOP/s"
sdpa_demo()
@part("6. Serialize a constructed graph, reload it, execute by UID")
def serialization():
Bsz, M, Kd, Nd = 8, 256, 512, 256
a = torch.randn(Bsz, M, Kd, system=DEV, dtype=DTYPE)
bm = torch.randn(Bsz, Kd, Nd, system=DEV, dtype=DTYPE)
out = torch.empty(Bsz, M, Nd, system=DEV, dtype=DTYPE)
UID_A, UID_B, UID_C = 1, 2, 3
g = cudnn.pygraph(
deal with=HANDLE, identify="serializable_mm",
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").set_uid(UID_A)
Bt = tensor_of(g, bm, "B").set_uid(UID_B)
C = g.matmul(A=A, B=Bt, compute_data_type=cudnn.data_type.FLOAT)
C.set_output(True).set_data_type(TORCH2CUDNN[DTYPE]).set_uid(UID_C)
t0 = time.perf_counter()
construct(g)
cold_ms = (time.perf_counter() - t0) * 1e3
blob = g.serialize()
print(f" chilly construct : {cold_ms:.1f} ms")
print(f" serialized plan : {len(blob)} bytes (cache this to disk / ship it)")
t0 = time.perf_counter()
g2 = cudnn.pygraph()
strive:
g2.deserialize(HANDLE, blob)
besides TypeError:
g2.deserialize(blob)
warm_ms = (time.perf_counter() - t0) * 1e3
print(f" deserialize : {warm_ms:.1f} ms -> {cold_ms/max(warm_ms,1e-6):.1f}x quicker startup")
ws = torch.empty(max(g2.get_workspace_size(), 1), system=DEV, dtype=torch.uint8)
g2.execute({UID_A: a, UID_B: bm, UID_C: out}, ws, deal with=HANDLE)
torch.cuda.synchronize()
ref = torch.bmm(a.float(), bm.float())
rel = ((out.float() - ref).abs().max() / ref.abs().max()).merchandise()
print(f" rel err after reload: {rel:.2e}")
return f"{len(blob)} B blob, reload {cold_ms/max(warm_ms,1e-6):.1f}x quicker than rebuild"
serialization()
We construct a fused scaled dot-product consideration graph with causal masking and verify it in opposition to torch.nn.practical.scaled_dot_product_attention, guarding the entire part behind an SM80 verify as a result of the fused kernels want Ampere or newer. We write the causal argument with fallbacks, for the reason that frontend has moved from use_causal_mask towards diagonal_alignment and certain arguments throughout its 1.x releases. We then serialize a constructed matmul graph to bytes, reload it right into a contemporary graph object, and execute it by way of integer UIDs, which lets us skip the compilation price solely at course of startup.
@part("7. Dynamic shapes with a shared kernel cache")
def dynamic_shapes():
kc = cudnn.create_kernel_cache()
def make(n):
x = torch.randn(n, 64, 32, 32, system=DEV, dtype=DTYPE).to(memory_format=torch.channels_last)
w = torch.randn(64, 64, 3, 3, system=DEV, dtype=DTYPE).to(memory_format=torch.channels_last)
y = torch.empty(n, 64, 32, 32, system=DEV, dtype=DTYPE).to(memory_format=torch.channels_last)
g = cudnn.pygraph(
deal with=HANDLE, identify=f"dyn_{n}",
io_data_type=TORCH2CUDNN[DTYPE],
intermediate_data_type=cudnn.data_type.FLOAT,
compute_data_type=cudnn.data_type.FLOAT,
kernel_cache=kc,
is_dynamic_shape_enabled=True,
)
X, Wt = tensor_of(g, x, "X"), tensor_of(g, w, "W")
Y = g.conv_fprop(picture=X, weight=Wt, padding=[1, 1], stride=[1, 1],
dilation=[1, 1], compute_data_type=cudnn.data_type.FLOAT)
Y.set_output(True).set_data_type(TORCH2CUDNN[DTYPE])
Y.set_dim(listing(y.dimension())).set_stride(listing(y.stride()))
t0 = time.perf_counter()
construct(g)
ms = (time.perf_counter() - t0) * 1e3
ws = workspace_for(g)
g.execute({X: x, Wt: w, Y: y}, ws)
torch.cuda.synchronize()
return ms
occasions = [(n, make(n)) for n in (8, 16, 24, 32)]
for n, ms in occasions:
print(f" batch {n:>3d}: construct {ms:7.1f} ms")
first, relaxation = occasions[0][1], [m for _, m in times[1:]]
print(f"n first form {first:.1f} ms, later shapes avg {sum(relaxation)/len(relaxation):.1f} ms")
print(" The cache lets shape-variant graphs reuse an already-JIT'd kernel,")
print(" which is what retains variable batch/seqlen serving out of rebuild hell.")
return f"first {first:.0f} ms vs subsequent {sum(relaxation)/len(relaxation):.0f} ms"
dynamic_shapes()
@part("8. CUDA Graph seize round a cuDNN execution plan")
def cuda_graph_capture():
if not CONV_STATE:
elevate RuntimeError("part 2 didn't run, nothing to seize")
g, pack, ws = CONV_STATE["graph"], CONV_STATE["pack"], CONV_STATE["ws"]
eager_ms = bench(lambda: g.execute(pack, ws))
aspect = torch.cuda.Stream()
aspect.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(aspect):
cudnn.set_stream(deal with=HANDLE, stream=aspect.cuda_stream)
for _ in vary(3):
g.execute(pack, ws, deal with=HANDLE)
torch.cuda.current_stream().wait_stream(aspect)
torch.cuda.synchronize()
cg = torch.cuda.CUDAGraph()
with torch.cuda.graph(cg):
cudnn.set_stream(deal with=HANDLE, stream=torch.cuda.current_stream().cuda_stream)
g.execute(pack, ws, deal with=HANDLE)
cudnn.set_stream(deal with=HANDLE, stream=torch.cuda.current_stream().cuda_stream)
replay_ms = bench(lambda: cg.replay())
report("plain execute()", eager_ms)
report("cuda graph replay()", replay_ms)
print(f" launch overhead eliminated: {(eager_ms-replay_ms)*1e3:.1f} us/iter")
print(" Pointers are frozen at seize time -- reuse the identical buffers and")
print(" copy new knowledge into them, or re-capture.")
return f"{eager_ms:.3f} -> {replay_ms:.3f} ms by way of replay"
cuda_graph_capture()
banner("SUMMARY")
for identify, res in RESULTS.objects():
print(f" {identify:<58s} {res}")
print("""
Where to go subsequent
- samples/python within the repo: FP8/MXFP8 consideration, paged KV cache, MoE grouped GEMM
- python/cudnn/: the open-sourced CuTe DSL kernels (SDPA, grouped GEMM + SwiGLU,
block-sparse and native sparse consideration) you'll be able to learn and modify
- debugging: CUDNN_FRONTEND_LOG_INFO=1 and CUDNN_FRONTEND_LOG_FILE=stdout
(use degree 10 throughout CUDA graph seize -- degree 1 dumps tensors and shouldn't be
capture-safe)
""")
We end with two manufacturing considerations. First, we share a kernel cache throughout 4 graphs that differ solely in batch dimension and time every construct, so we will see later shapes reuse an already compiled kernel as an alternative of paying the JIT price once more. Then we seize the convolution plan inside a CUDA graph, setting the cuDNN deal with’s stream to the seize stream. So the work lands within the graph, and we measure how a lot per-iteration launch overhead the replay removes.
In conclusion, what we constructed right here was small in code however broad in scope: a convolution, a matmul, and an consideration kernel, every expressed as a graph reasonably than a library name. Working at that degree modified what we may determine. We selected which operations collapsed right into a single kernel, so the bias provides, activations, and AMAX reductions we folded into the epilogues by no means wrote an intermediate to reminiscence. We selected the engine ourselves as an alternative of accepting a heuristic, and timing each candidate config advised us what that selection was price. We additionally selected when to pay for compilation, pushing it out of the recent path with serialized plans, a kernel cache shared throughout shapes, and CUDA graph seize. The checks in opposition to PyTorch mattered as a lot because the timings, for the reason that locations the place we merely matched it have been normally locations the place PyTorch was already calling cuDNN beneath. That marked out the place this API earns its hold: fusions with no framework-level equal, shapes sizzling sufficient to justify autotuning, and small kernels the place startup and launch prices dominate.
Check out the FULL CODES here. All credit score goes to the researcher of this mission. Also, be happy to comply with us on Twitter and don’t overlook to hitch our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
Need to companion with us for selling your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar and so on.? Connect with us
The submit Inside NVIDIA’s cuDNN Graph API: Fusion, Autotuning, and Plan Reuse with cuDNN Frontend appeared first on MarkTechPost.
