cuDNN: add fused scaled-dot-product (flash) attention forward by CarloLucibello · Pull Request #3174 · JuliaGPU/CUDA.jl

CarloLucibello · 2026-06-11T13:21:32Z

Adds a fused scaled-dot-product (flash) attention forward pass to the cuDNN extension, built on cuDNN's modern backend graph API.

What this adds

lib/cudnn/src/backend.jl — a thin Julia wrapper over the cuDNN backend graph API (cudnnBackend*): typed descriptor wrapper, setattr!/getattr, bfinalize!, and helpers for tensors, operation graphs, engine heuristics, execution plans, and variant packs.

lib/cudnn/src/sdpa.jl — cudnnSDPAForward[!] using CUDNN_BACKEND_OPERATION_SDPA_FWD_DESCRIPTOR (the dedicated single-node fused attention op):

cudnnSDPAForward(q, k, v; scale=1/√d)   # → out
cudnnSDPAForward!(out, q, k, v; scale)

Input/output layout is (head_dim, num_heads, seq_len, batch) — matches NNlib's internal multi-head attention layout directly, no permute needed. Plans are cached per (device, dtype, d, sq, skv, h, b).

Supported / not yet supported

	Status
Float16, BFloat16	✓
Ampere (sm_80) or newer	✓
Forward inference	✓
Causal masking	blocked on cuDNN ≤ 9.20 (requires score-modifier subgraph, cuDNN ≥ 9.21)
Backward	blocked on cuDNN ≤ 9.20 (same limitation)

Float32/Float64 are rejected with a clear error (not supported by the fused engine).

BFloat16 fix

An early version had a hang on BFloat16 caused by GC finalizers (cudnnBackendDestroyDescriptor) firing inside a @gcsafe_ccall to cudnnBackendExecute during cuDNN's JIT compilation of runtime-compiled engines, contending for cuDNN's internal state lock. Fixed by destroying unused descriptors synchronously in getattr_descriptors and try_execution_plan rather than leaving GC-pending finalizers. Both Float16 and BFloat16 are now tested.

Tested on

RTX 5090 (sm_120), cuDNN 9.20, Julia 1.12.

CarloLucibello · 2026-06-17T08:19:39Z

runners are executing tests with missing cuDNN runtime
https://buildkite.com/julialang/cuda-dot-jl/builds/7491/canvas?jid=019ed48f-4aff-4045-8290-4724609610cf&tab=output#L476

codecov · 2026-06-17T08:40:20Z

Codecov Report

❌ Patch coverage is 0% with 177 lines in your changes missing coverage. Please review.
✅ Project coverage is 17.16%. Comparing base (f36a98c) to head (815104f).

Files with missing lines	Patch %	Lines
lib/cudnn/src/backend.jl	0.00%	90 Missing ⚠️
lib/cudnn/src/sdpa.jl	0.00%	62 Missing ⚠️
lib/cudnn/test/sdpa.jl	0.00%	25 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3174      +/-   ##
==========================================
+ Coverage   16.88%   17.16%   +0.27%     
==========================================
  Files         124      127       +3     
  Lines        9886    10063     +177     
==========================================
+ Hits         1669     1727      +58     
- Misses       8217     8336     +119

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

maleadt · 2026-06-17T11:56:49Z

runners are executing tests with missing cuDNN runtime

That's because you're only running the Julia tests, not the CUDA ones. Those are skipped for draft PRs, and your PR was marked as such at the time of the last commit.

github-actions

CUDA.jl Benchmarks

Details

Benchmark suite	Current: `e7db729`	Previous: `1b62bc2`	Ratio
`array/accumulate/Float32/1d`	`97873` ns	`98427` ns	`0.99`
`array/accumulate/Float32/dims=1`	`74125` ns	`73956` ns	`1.00`
`array/accumulate/Float32/dims=1L`	`1599310` ns	`1599743` ns	`1.00`
`array/accumulate/Float32/dims=2`	`139386` ns	`139832` ns	`1.00`
`array/accumulate/Float32/dims=2L`	`659136` ns	`659439` ns	`1.00`
`array/accumulate/Int64/1d`	`117753` ns	`117262` ns	`1.00`
`array/accumulate/Int64/dims=1`	`77723` ns	`77754` ns	`1.00`
`array/accumulate/Int64/dims=1L`	`1714110` ns	`1716436` ns	`1.00`
`array/accumulate/Int64/dims=2`	`151015` ns	`150712` ns	`1.00`
`array/accumulate/Int64/dims=2L`	`985888` ns	`986611` ns	`1.00`
`array/broadcast`	`17869` ns	`18091` ns	`0.99`
`array/construct`	`845.1041666666666` ns	`839.3695652173913` ns	`1.01`
`array/copy`	`15900` ns	`16125` ns	`0.99`
`array/copyto!/cpu_to_gpu`	`206605` ns	`208361` ns	`0.99`
`array/copyto!/gpu_to_cpu`	`240379` ns	`241050` ns	`1.00`
`array/copyto!/gpu_to_gpu`	`8826.666666666666` ns	`8767.666666666666` ns	`1.01`
`array/iteration/findall/bool`	`131022` ns	`131721` ns	`0.99`
`array/iteration/findall/int`	`145328` ns	`145411` ns	`1.00`
`array/iteration/findfirst/bool`	`67237` ns	`67968` ns	`0.99`
`array/iteration/findfirst/int`	`68791` ns	`68551` ns	`1.00`
`array/iteration/findmin/1d`	`60260` ns	`60753` ns	`0.99`
`array/iteration/findmin/2d`	`98841` ns	`99184` ns	`1.00`
`array/iteration/logical`	`183702` ns	`184742` ns	`0.99`
`array/iteration/scalar`	`60847` ns	`61604` ns	`0.99`
`array/permutedims/2d`	`48525` ns	`48136` ns	`1.01`
`array/permutedims/3d`	`49706` ns	`50151` ns	`0.99`
`array/permutedims/4d`	`49288` ns	`49687` ns	`0.99`
`array/random/rand/Float32`	`11644` ns	`11835` ns	`0.98`
`array/random/rand/Int64`	`22864` ns	`23179` ns	`0.99`
`array/random/rand!/Float32`	`7648.75` ns	`7672.25` ns	`1.00`
`array/random/rand!/Int64`	`17984` ns	`17902` ns	`1.00`
`array/random/randn/Float32`	`32499` ns	`32829` ns	`0.99`
`array/random/randn!/Float32`	`23079` ns	`23288` ns	`0.99`
`array/reductions/mapreduce/Float32/1d`	`30514` ns	`30921` ns	`0.99`
`array/reductions/mapreduce/Float32/dims=1`	`36864` ns	`37530` ns	`0.98`
`array/reductions/mapreduce/Float32/dims=1L`	`49789` ns	`49965` ns	`1.00`
`array/reductions/mapreduce/Float32/dims=2`	`54289` ns	`54519` ns	`1.00`
`array/reductions/mapreduce/Float32/dims=2L`	`65661` ns	`65926` ns	`1.00`
`array/reductions/mapreduce/Int64/1d`	`37503` ns	`37843` ns	`0.99`
`array/reductions/mapreduce/Int64/dims=1`	`39755` ns	`39986` ns	`0.99`
`array/reductions/mapreduce/Int64/dims=1L`	`87831` ns	`88012` ns	`1.00`
`array/reductions/mapreduce/Int64/dims=2`	`56663` ns	`57110` ns	`0.99`
`array/reductions/mapreduce/Int64/dims=2L`	`81751` ns	`81936` ns	`1.00`
`array/reductions/reduce/Float32/1d`	`30813` ns	`31404` ns	`0.98`
`array/reductions/reduce/Float32/dims=1`	`36910` ns	`37293` ns	`0.99`
`array/reductions/reduce/Float32/dims=1L`	`49709` ns	`49872` ns	`1.00`
`array/reductions/reduce/Float32/dims=2`	`54275` ns	`54508` ns	`1.00`
`array/reductions/reduce/Float32/dims=2L`	`67253` ns	`67427` ns	`1.00`
`array/reductions/reduce/Int64/1d`	`38119` ns	`37603` ns	`1.01`
`array/reductions/reduce/Int64/dims=1`	`39139` ns	`40000` ns	`0.98`
`array/reductions/reduce/Int64/dims=1L`	`87842` ns	`87926` ns	`1.00`
`array/reductions/reduce/Int64/dims=2`	`56340` ns	`56940` ns	`0.99`
`array/reductions/reduce/Int64/dims=2L`	`81426` ns	`82266` ns	`0.99`
`array/reverse/1d`	`16003` ns	`16174` ns	`0.99`
`array/reverse/1dL`	`68810` ns	`68908` ns	`1.00`
`array/reverse/1dL_inplace`	`66910` ns	`66896` ns	`1.00`
`array/reverse/1d_inplace`	`8499.333333333334` ns	`8113` ns	`1.05`
`array/reverse/2d`	`19005` ns	`19235` ns	`0.99`
`array/reverse/2dL`	`72457` ns	`72424` ns	`1.00`
`array/reverse/2dL_inplace`	`66611` ns	`66960` ns	`0.99`
`array/reverse/2d_inplace`	`9962` ns	`9625` ns	`1.04`
`array/sorting/1d`	`2654465` ns	`2667955` ns	`0.99`
`array/sorting/2d`	`1038499` ns	`1037928` ns	`1.00`
`array/sorting/by`	`3191863` ns	`3192120` ns	`1.00`
`cuda/synchronization/context/auto`	`1052.7` ns	`1052.6` ns	`1.00`
`cuda/synchronization/context/blocking`	`806.4891304347826` ns	`791.8888888888889` ns	`1.02`
`cuda/synchronization/context/nonblocking`	`5698` ns	`5760.833333333333` ns	`0.99`
`cuda/synchronization/stream/auto`	`874.9803921568628` ns	`902.6904761904761` ns	`0.97`
`cuda/synchronization/stream/blocking`	`677.4285714285714` ns	`674.8039215686274` ns	`1.00`
`cuda/synchronization/stream/nonblocking`	`5552.666666666667` ns	`5552` ns	`1.00`
`integration/byval/reference`	`147265` ns	`147376` ns	`1.00`
`integration/byval/slices=1`	`149256` ns	`149419` ns	`1.00`
`integration/byval/slices=2`	`291856` ns	`292077` ns	`1.00`
`integration/byval/slices=3`	`434499` ns	`434464` ns	`1.00`
`integration/cudadevrt`	`104361` ns	`104457` ns	`1.00`
`integration/volumerhs`	`9305353` ns	`9296159` ns	`1.00`
`kernel/indexing`	`12289` ns	`12401` ns	`0.99`
`kernel/indexing_checked`	`13036` ns	`13288` ns	`0.98`
`kernel/launch`	`1897.1` ns	`1989.9` ns	`0.95`
`kernel/occupancy`	`645.5421686746988` ns	`641.1878787878788` ns	`1.01`
`kernel/rand`	`15373` ns	`14274` ns	`1.08`
`latency/import`	`3911506976` ns	`3917545175` ns	`1.00`
`latency/precompile`	`4673782976` ns	`4670383446` ns	`1.00`
`latency/ttfp`	`4894922227` ns	`4930163586` ns	`0.99`

This comment was automatically generated by workflow using github-action-benchmark.

CarloLucibello · 2026-06-18T18:48:40Z

+# maxcount-n slots: if a finalizer fires cudnnBackendDestroyDescriptor while another thread
+# is inside a cudnnBackendExecute that holds cuDNN's JIT lock, we get a deadlock on bf16
+# and other runtime-compiled engines. Instead we use raw handles, destroy the unused ones
+# synchronously right here, and only register finalizers for the n handles we return.


@maleadt is this manual memory management ok?

CarloLucibello · 2026-06-18T18:49:06Z

this is ready for review/merge

Wrap cuDNN's modern fused SDPA via the backend graph API. The legacy `cudnnMultiHeadAttnForward` is the deprecated attention path; this adds the flash-attention kernel NVIDIA now recommends. - backend.jl: a thin typed layer over the cuDNN backend graph API (`cudnnBackend*`): descriptor wrapper, setattr!/getattr, and helpers to build an operation graph, run engine heuristics, finalize an execution plan, and execute a variant pack. No prior high-level wrapper used this API. - sdpa.jl: `cudnnSDPAForward[!]` driving the dedicated `CUDNN_BACKEND_OPERATION_SDPA_FWD_DESCRIPTOR`, with a per-shape execution-plan cache. Inputs are 4-D (head_dim, nheads, seq_len, batch) — NNlib's attention layout — so no permute is needed to interoperate. Scope: forward inference only, Float16/BFloat16 (cuDNN's fused engine does not support Float32/Float64). Verified vs a Float32 reference (relerr ~5e-4). Not yet supported, both blocked on cuDNN <= 9.20 and documented inline: - causal masking: the SDPA score-modifier subgraph needs cuDNN >= 9.21 (no CUDNN_jll yet); block-mask is block-sparse; the primitive matmul->softmax ->matmul graph yields no fused engine from raw backend calls. - backward: `cudnnSDPABackward` is a documented placeholder; the dedicated SDPA_BWD descriptor does not finalize on 9.20 and the supported path needs the same 9.21 subgraph mechanism (forward stats output already verified to work). Tests cover the forward against a dense reference (Float16, several shapes, custom scale) plus in-place agreement; gated on compute capability >= 8.0. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

- Constrain cudnnSDPAForward[!] inputs to DenseCuArray{T,4}: the cached plan bakes in dense column-major strides, so non-contiguous views (or host Arrays, whose Ptr/CuPtr distinction the variant pack erases) previously passed all size checks and produced silently wrong results or illegal accesses. - Key the plan cache on the current context: execution plans are finalized against a specific device's handle, so a plan built on one GPU must not be executed on another. - Only swallow the CUDNN_STATUS_NOT_SUPPORTED family (3000s) in try_execution_plan; BAD_PARAM/INTERNAL_ERROR now propagate instead of being misreported as "no supported engine". The terminal error also distinguishes "heuristic returned no configs" from "N configs failed to finalize". - Build plans outside the cache lock (matching the descriptors.jl pattern) so concurrent calls don't serialize behind a multi-millisecond plan build; a racing duplicate build is benign and resolved by get! on insert. - Use with_workspace for the execute workspace, freeing it eagerly instead of leaving it to GC (house style, cf. convolution.jl/reduce.jl). - Test that views, host arrays, and Float32 inputs are rejected. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Fix a GC-finalizer/cuDNN-JIT deadlock that caused BFloat16 fused SDPA to hang indefinitely on runtime-compiled engines (e.g. Blackwell sm_120): - getattr_descriptors: use raw handles internally, destroy unused slots synchronously before any long operation runs, register finalizers only for the n returned descriptors. Previously, up to 12 unreachable wrapper objects with cudnnBackendDestroyDescriptor finalizers were left pending; since all backend calls use @gcsafe_ccall, the GC could fire those finalizers during a JIT-compiling cudnnBackendExecute, contending for cuDNN's internal state lock -> deadlock. - try_execution_plan: destroy the descriptor immediately on failure instead of leaving a pending finalizer. - _build_sdpa_plan: keep ALL engine configs alive in SDPAPlan.keepalive (not just the first successful one), preventing the unused cfgs from becoming unreachable between plan build and execute. Also enable BFloat16 in the SDPA tests: switch to CUDA.randn for test data (GPU-side generation avoids the CPU BFloat16.(randn(Float32,...)) broadcast which triggers a slow first-time LLVM compilation on some machines), and add BFloat16 alongside Float16 in the test loop. Tested on RTX 5090 (sm_120), cuDNN 9.20: Float16 and BFloat16 both pass all 5 shape variants × 4 assertions = 40 tests, plus the 3 invalid-input rejection tests. Remove scratch investigation files (sdpa_bf16_findings.md, sdpa_bf16_hang.jl, torch_sdpa_test.py). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Use an idempotent backend descriptor destroy path and a Julia-level finalize! helper for descriptor finalization. Add device-property descriptors to backend heuristics and execution plans, align the SDPA causal keyword with common SDPA APIs, and cover the backend wrapper plus SDPA test data generation through cuRAND.

CarloLucibello marked this pull request as draft June 11, 2026 13:21

CarloLucibello force-pushed the cl/cudnn-sdpa branch 2 times, most recently from 65ed835 to b75073c Compare June 17, 2026 07:49

CarloLucibello marked this pull request as ready for review June 17, 2026 07:54

CarloLucibello closed this Jun 17, 2026

CarloLucibello reopened this Jun 17, 2026

github-actions Bot reviewed Jun 17, 2026

View reviewed changes

CarloLucibello commented Jun 18, 2026

View reviewed changes

CarloLucibello and others added 7 commits July 9, 2026 08:15

cuDNN: fix dropout state dict keying

4c85c84

cuDNN: tighten backend descriptor lifetimes

2fff907

cuDNN: pool SDPA plans with handles

e7db729

maleadt force-pushed the cl/cudnn-sdpa branch from 815104f to e7db729 Compare July 9, 2026 07:23

maleadt mentioned this pull request Jul 9, 2026

Modernize cuDNN wrappers around the backend graph API #3191

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

cuDNN: add fused scaled-dot-product (flash) attention forward#3174

cuDNN: add fused scaled-dot-product (flash) attention forward#3174
CarloLucibello wants to merge 7 commits into
JuliaGPU:mainfrom
CarloLucibello:cl/cudnn-sdpa

CarloLucibello commented Jun 11, 2026 •

edited

Loading

Uh oh!

CarloLucibello commented Jun 17, 2026

Uh oh!

codecov Bot commented Jun 17, 2026 •

edited

Loading

Uh oh!

maleadt commented Jun 17, 2026

Uh oh!

github-actions Bot left a comment •

edited

Loading

Uh oh!

CarloLucibello Jun 18, 2026

Uh oh!

CarloLucibello commented Jun 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

CarloLucibello commented Jun 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What this adds

Supported / not yet supported

BFloat16 fix

Tested on

Uh oh!

CarloLucibello commented Jun 17, 2026

Uh oh!

codecov Bot commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

maleadt commented Jun 17, 2026

Uh oh!

github-actions Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

CUDA.jl Benchmarks

Uh oh!

CarloLucibello Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

CarloLucibello commented Jun 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

CarloLucibello commented Jun 11, 2026 •

edited

Loading

codecov Bot commented Jun 17, 2026 •

edited

Loading

github-actions Bot left a comment •

edited

Loading