@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, k, v, is_causal=causal, scale=scale)
rel = ((o.float() - ref.float()).abs().max() / ref.float().abs().max()).item()
print(f" shape : 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.functional.scaled_dot_product_attention(
q, k, v, is_causal=causal, scale=scale))
print()
report("cuDNN FE SDPA", ms, SDPA_FLOPS)
report("torch SDPA (backend's choice)", ms_t, SDPA_FLOPS)
print(" Note: torch may already be dispatching to cuDNN or FlashAttention,")
print(" so parity here is the expected, healthy outcome.")
return f"{ms:.3f} ms, {tflops(SDPA_FLOPS, ms):.1f} TFLOP/s"
sdpa_demo()
@section("6. Serialize a built graph, reload it, execute by UID")
def serialization():
Bsz, M, Kd, Nd = 8, 256, 512, 256
a = torch.randn(Bsz, M, Kd, device=DEV, dtype=DTYPE)
bm = torch.randn(Bsz, Kd, Nd, device=DEV, dtype=DTYPE)
out = torch.empty(Bsz, M, Nd, device=DEV, dtype=DTYPE)
UID_A, UID_B, UID_C = 1, 2, 3
g = cudnn.pygraph(
handle=HANDLE, name="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()
build(g)
cold_ms = (time.perf_counter() - t0) * 1e3
blob = g.serialize()
print(f" cold build : {cold_ms:.1f} ms")
print(f" serialized plan : {len(blob)} bytes (cache this to disk / ship it)")
t0 = time.perf_counter()
g2 = cudnn.pygraph()
try:
g2.deserialize(HANDLE, blob)
except 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 faster startup")
ws = torch.empty(max(g2.get_workspace_size(), 1), device=DEV, dtype=torch.uint8)
g2.execute({UID_A: a, UID_B: bm, UID_C: out}, ws, handle=HANDLE)
torch.cuda.synchronize()
ref = torch.bmm(a.float(), bm.float())
rel = ((out.float() - ref).abs().max() / ref.abs().max()).item()
print(f" rel err after reload: {rel:.2e}")
return f"{len(blob)} B blob, reload {cold_ms/max(warm_ms,1e-6):.1f}x faster than rebuild"
serialization()
