perf(cuda): make the cuDNN convolution path usable, and 4.4x faster by jamesbrink · Pull Request #11 · utensils/candle

jamesbrink · 2026-08-31T03:59:14Z

Brings the fork's main up to what mold already ships, and makes the cudnn feature worth compiling.

Why this is 5 commits, not 1

main is behind the revision mold currently pins (5de41be7). The first three commits are already in production use through that pin; cherry-picking only the cuDNN work onto main would produce a revision missing the ConvRot and neighborhood-attention kernels mold depends on, and would not build. So this catches main up rather than landing cuDNN in isolation.

commit	what
`cd0df80b`	INT8 ConvRot dequantize kernel
`4bf88883`	neighborhood-attention-3d CUDA kernel
`5de41be7`	NA3D degrades to full attention when the kernel exceeds grid extent
`bcb30d43`	cuDNN convolution dispatch, math types, workspace
`2a157323`	per-thread policy, honest dispatch counting, test fixes

The cuDNN work

candle had a cudnn feature that appeared only in READMEs — no tests, and every convolution descriptor left at CUDNN_DEFAULT_MATH. That declines tensor-core engines for f16/bf16 (a bf16 convolution ran slower than the same convolution in f32) and silently permits TF32 for f32. Selecting TENSOR_OP for half precision and FMA for wide dtypes is 2.4x on bf16 at bit-identical output, and on f32 it is both 1.4x faster and ~100x more accurate. candle also zeroed the cuDNN workspace on every call — a memset of up to hundreds of MB that cuDNN never reads.

Measured on an RTX 4090, one steady-state latent frame of a Wan 2.1 VAE decoder at 832x480 in bf16, summing every convolution:

path	per latent frame	vs im2col
im2col	845 ms	-
cuDNN as shipped	518 ms	1.63x
cuDNN after this	192 ms	4.38x

cuDNN is not always right, so dispatch is per convolution. A 1x1 kernel is a matmul whose im2col buffer is a free reshape, and a small convolution finishes in less time than cuDNN spends creating descriptors and running its algorithm heuristic — neither of which candle caches. The gate is the size of the column buffer im2col would materialise, which is exactly the work cuDNN avoids. Output elements are the wrong measure: this decoder's head convolution (96->3 at 480x832) has barely a million output elements, a 690 MB column buffer, and runs 6.9x faster on cuDNN. Any cuDNN error falls back to im2col rather than failing the call.

Policy

cudnn_policy is a per-thread switch. A process-global would be wrong for any host driving several devices — those run one worker thread per device, and a global lets a thread that wants cuDNN flip it underneath a thread mid-convolution that needs the reproducible path, silently. A CUDA op executes inline on the issuing thread, so thread-local makes "the policy this thread set" and "the policy this convolution ran under" the same thing. Default is im2col: a caller that expressed no preference gets the path whose bytes do not move.

A pre-existing correctness fix

The im2col non-contiguous-kernel branch matmul'd the original strided kernel instead of the packed kernel_c. conv2d_grad_noncontiguous_kernel_gpu fails on main with --features cuda and passes with cuda,cudnn only because cuDNN masked it. Fixed at both CUDA sites. CPU and Metal have the identical pattern; the CPU test passes and I have no Metal hardware, so both are left alone.

Tests

candle-core/tests/cudnn_tests.rs is new — the feature had none. Every comparison asserts via cudnn_policy::dispatch_count() that a convolution actually executed on cuDNN: dispatch is conditional several times over, so a suite that only sets the policy compares im2col to itself and passes no matter how broken cuDNN is. An earlier draft of this file did exactly that — mutating the cuDNN padding left all 8 tests green.

Verified: 9/9 pass on an RTX 4090 (cuDNN 9.13); reverting the math type fails 5 of them; mutating the padding fails 6. clippy clean with -D warnings on default features and on cuda,cudnn (lib + tests — the two bench_main findings predate this branch).

🤖 Generated with Claude Code

https://claude.ai/code/session_01AjPWQj3ZabxXoMRhptD4c9

`dequantize_int8_convrot_256` had a CPU and a Metal arm only, so every CUDA consumer of a Comfy int8-convrot checkpoint (mold's LTX-2.5 packs and their Gemma 4 encoder) decoded each weight on the host and uploaded the widened BF16 tensor. This adds `candle-kernels/src/convrot.cu`, a line-for-line port of the Metal kernel (one 256-thread block per row group, the same three-round Hadamard butterfly in shared memory, one power-of-two scale multiply, one RNE narrowing), registers it as the `CONVROT` module, and wires `cuda_fwd` on the custom op. The CUDA arm is bit-identical to the CPU op, pinned by `cuda_matches_cpu_for_tiny_bounded_inputs`; the 16384x4096 residency smoke test mirrors the Metal one. Claude-Session: https://claude.ai/code/session_01DQgGe9N4ipofzmQ91SSHNT

Line-for-line port of neighborhood_attention.metal: one 256-thread block per (batch, query-position, head) computes the windowed attention score/softmax/accumulate in three phases over dynamic shared memory, dispatched through Map3 for F32/F16/BF16. Guards against oversized kernels before dispatch (48 KiB default CUDA per-block shared memory budget) the same way the Metal backend checks maxThreadgroupMemoryLength. Unblocks LTX-2.5's plain-bf16 diffusion-VAE decode on CUDA.

…ernel exceeds grid extent validate() unconditionally refused any kernel axis wider than its grid extent — a check shared by all three backends, but only ever reached by CPU/Metal in practice since CUDA previously failed earlier with "no cuda implementation". LTX-2.5's diffusion-VAE decoder uses a fixed [11, 11, 11] window regardless of clip length, so a short clip (e.g. a 9-frame smoke render) always hits this on every backend. Per NATTEN boundary-window semantics, a window wider than the axis it slides over degrades to full attention on that axis (start = 0, width = extent) rather than refusing. Add effective_kernel() to clamp each axis to min(kernel, extent) and thread it through cpu_fwd, cuda_fwd, and metal_fwd in place of the raw request kernel; validate() now only checks that the *requested* kernel is a positive odd triple. Add CUDA-vs-CPU parity coverage at the exact repro shape (grid [9, 64, 64], kernel [11, 11, 11]) and a multi-axis edge case where every axis clamps at once, including two even extents ([1, 4, 4] grid, kernel [3, 7, 7]).

candle has shipped a complete cuDNN conv1d/conv2d implementation behind `--features cudnn` with no test coverage, no runtime control, and a math type that defeats the point of using cuDNN at all. Measured on an RTX 4090 (cuDNN 9.13) against the Wan 2.1 VAE decoder's convolution shapes, one steady-state latent frame at 832x480, bf16: im2col (candle's default) 845 ms cudnn, as it was 518 ms 1.63x cudnn, after this commit 192 ms 4.38x Four changes. **Ask cuDNN for the right math mode.** Descriptors default to `CUDNN_DEFAULT_MATH`, which declines tensor-core engines for f16/bf16 — so a bf16 convolution ran *slower* than the same convolution in f32 — and silently admits TF32 for f32, truncating the mantissa to 10 bits. Selecting `TENSOR_OP` for half and `FMA` for wide dtypes fixes both, and neither is a trade: TENSOR_OP is 2.4x faster on bf16 with bit-identical output, and FMA is 1.4x faster on f32 *and* 100x more accurate (3.2e-6 vs 3.3e-4 peak relative deviation from im2col). **Stop zeroing the workspace.** cuDNN treats it as scratch and never reads before writing, so `alloc_zeros` was memsetting up to hundreds of MB per call. Worth ~10% of the decode's convolution time. **Choose per convolution, not per build.** The two implementations were `#[cfg]`-exclusive, so cuDNN was all-or-nothing at compile time and a cuDNN error killed the operation outright with im2col sitting right there. Now both compile, `cudnn_policy` gates the path at runtime, any cuDNN error falls back to im2col, and shapes where cuDNN's uncached per-call descriptor and algorithm setup costs more than it saves stay on im2col. The threshold is the size of the column buffer im2col materialises, which is the work cuDNN actually avoids; output elements are the wrong measure, as the Wan head convolution shows — 96->3 at 480x832 has barely a million output elements, a 690 MB column buffer, and runs 6.9x faster on cuDNN. The runtime switch also exists so an application can keep one model family bit-reproducible while another takes the speedup; the two paths sum in a different order and never agree exactly. **Fix the non-contiguous kernel path.** Both im2col convolutions packed the kernel into a contiguous `kernel_c` and then matmul'd the *original* strided buffer, described with a contiguous layout and the original's start offset — reading the wrong elements. `conv2d_grad_noncontiguous_kernel_gpu` fails on `--features cuda` before this commit and passes after. The identical pattern is present in the CPU and Metal backends and is left alone here: the CPU test passes and no Metal hardware was available to verify a change. Tests: `candle-core/tests/cudnn_tests.rs`, run on the 4090. Each parity test asserts via `cudnn_policy::dispatch_count()` that a convolution really reached cuDNN — dispatch is conditional three times over, so a suite that merely enables the policy compares im2col to itself and passes however broken cuDNN is. An earlier draft of this file did exactly that: mutating the cuDNN padding left all of it green. The committed version fails 6 of 9 on that mutation and 4 of 9 on reverting the math type. Full candle-core suite: 349 passed, 0 failed.

Peer review found three defects in the policy plumbing and three in the tests meant to guard it. The switch was a process-global AtomicBool, justified by "the host runs one model at a time". That is not true of a host driving several devices: those run one worker thread per device, and a global lets a thread that wants cuDNN flip it underneath a thread mid-convolution that needs the reproducible path. The failure is silent — the wrong-path convolution still returns a plausible answer, just not the bytes that were promised. A CUDA op executes inline on the issuing thread, so thread-local state makes "the policy this thread set" and "the policy this convolution ran under" the same thing by construction. The default is now im2col rather than cuDNN, so a caller that expressed no preference gets the path whose bytes do not move. Threads do not inherit it, so no worker can perturb another. Two tests pin the isolation. record_dispatch fired at the dispatch decision, but the cuDNN arm returns Ok even when it handed a non-contiguous kernel to Candle's own Conv2D kernel — so the counter reported work cuDNN never did, to the very parity tests that gate on it. It now fires inside launch_conv2d and launch_conv1d, after the last point a fallback is possible. The counter is per-thread for the same reason the flag is, so concurrent tests each observe only their own convolutions and the suite needs no mutex. record_dispatch is now #[cfg(feature = "cudnn")]: with the feature off every caller compiles out and the helper tripped dead_code under the repository's -D warnings clippy job. On the tests: `cuda` without `cudnn` is a supported configuration where the device opens and is_enabled() is permanently false, so every comparison saw zero dispatches and failed for a reason that is not a defect; the suite now takes its device from cudnn_device(), which declines when the feature is absent. bf16 convolution needs SM80+, and on an older card cuDNN cannot execute it, so cudnn_runs_dtype() probes once with a shape known to clear the threshold — separating "this card cannot" from "the dispatch decision is wrong", which must still fail. And the strided-input test was not strided: it narrowed only the batch dimension, and Candle ignores strides on length-1 dimensions, so the input stayed is_contiguous() and never reached create_4d_tensor_ex. It now narrows width and asserts !is_contiguous() so it cannot silently stop testing that path again. The below-threshold case tuple is named, which tripped clippy::type_complexity under -D warnings. Verified: clippy clean on default features and on cuda,cudnn (lib and tests; the two bench findings predate this branch). 9/9 tests pass, and reverting the math type still fails 5 of them — one more than before, because the strided test now actually strides. Claude-Session: https://claude.ai/code/session_01AjPWQj3ZabxXoMRhptD4c9

Carried in unformatted from the NA3D commits this PR brings to main; the fork's Rustfmt CI job fails on it. Claude-Session: https://claude.ai/code/session_01AjPWQj3ZabxXoMRhptD4c9

utensils/candle#11 landed the cuDNN convolution path on the fork's main as 744ae3b8. Re-pins all five sites, both lockfiles and the Nix vendored source hash to that commit, so the branch depends on a revision that is actually on main rather than on a feature branch that no longer exists. candle-single-identity and desktop-candle-nix-source-hash both pass. Claude-Session: https://claude.ai/code/session_01AjPWQj3ZabxXoMRhptD4c9

…#1485) Closes #1483. Wan's VAE decode is the largest phase of a render outside the denoise, and ~76% of it is convolution. candle ran every one of those through im2col + GEMM, building a column buffer up to 690 MB per call. ## Measured, not projected Full renders on an RTX 4090, `wan22-t2v-a14b:q5`, 832x480, 81 frames, 4 steps, seed 1483, against the production model store. Both arms warm and parking the same blocks: | `MOLD_CONV` | Decode | Render | | ----------- | -----: | -----: | | `im2col` | 23.3 s | 105.9 s | | `cudnn` | **11.2 s** | **94.9 s** | | | **2.08x** | **1.12x** | Decode saves 12.1 s and the render saves 11.0 s — the gain lands entirely in the phase it should, which is what makes the attribution credible. im2col decode reproduced at 23.4/23.3 s across runs, so run-to-run noise is ~0.1 s. The first A/B looked like a 32.5 s win. It wasn't: the control had read the 14 GB checkpoint cold while the treatment found it in page cache. Re-running the control warm gave the honest 11.0 s. Frames match at ~40 dB PSNR across three sampled frames and are visually identical — same shot, not the same bytes, which is the trade this family opts into. Underneath, the convolutions themselves are **4.38x** (845 -> 192 ms per latent frame). A 4.4x on convolutions is only a 2.08x on the decode because a quarter of the decode is not convolution and cuDNN pays a per-call setup candle does not cache. The render number is the honest one to quote. ## Why enabling cuDNN alone was only a third of it candle created every convolution descriptor at `CUDNN_DEFAULT_MATH`, which declines tensor-core engines for f16/bf16 and silently permits TF32 for f32. Fixed in the fork (utensils/candle#11, merged): `TENSOR_OP` for half precision, `FMA` for wide dtypes — 2.4x on bf16 at bit-identical output, and on f32 both 1.4x faster and ~100x more accurate. Stock cuDNN was 1.63x; tuned is 4.38x. ## Scope: video families only `ConvPolicy` mirrors `AttentionPolicy` on the identical clips-versus-stills argument (#736), and a unit test asserts the two family lists never diverge. Image families stay on im2col in every build, so an archived still seed renders the same bytes forever. `MOLD_CONV={cudnn,im2col}` overrides either direction and is registered as engine-shaping. The policy is **per thread**, not per process. mold spawns one GPU worker thread per device, so a process-global switch would let a video render flip cuDNN on while an image render was mid-convolution on another device — silently moving still bytes. A CUDA op executes inline on the issuing thread, so thread-local makes the policy and the execution the same thing by construction. ## Peer review found four defects in this branch `codex review` on both repos returned 11 findings; all are fixed. The three that mattered: - **The scope missed chain stages.** It lived in `InferenceEngine::generate`, but a chain stage enters via `ChainStageRenderer::render_stage`. Since `mold run --frames 200` renders long video *as a chain*, the clips with the most to gain would have got nothing and `MOLD_CONV` would have been ignored for them. The 81-frame UAT is single-shot, so it never touched that path. - **The resolved backend was missing from the execution identity.** One source revision built with and without `cudnn` produced a single fingerprint for two different sets of numerics, so sharing a `mold.db` reused the other's scheduler estimates. Now carried in `ExecutionSemanticConfig`, skipped when `None` so image families' fingerprints are byte-identical to before — the technique `umt5_variant` already uses. - **The dispatch counter over-counted.** The cuDNN arm returns `Ok` even when it hands a non-contiguous kernel to candle's own kernel, so the counter reported work cuDNN never did — to the very parity tests that gate on it. After the thread-local rewrite I re-ran the UAT rather than trusting the unit tests: **11.1 s decode**, so the policy still dispatches. Had the scope and the convolutions landed on different threads, every test would still have passed and cuDNN would silently never have fired. ## Release artifacts Previously only `flake.nix` enabled the feature, while `.claude/rules/inference.md` claimed every Linux CUDA recipe did. Now all of them: `release.yml` (4 binaries), the Dockerfile (builder **and** runtime — the binary links `libcudnn.so.9` and the builder's own ldd check cannot catch its absence), both AUR PKGBUILDs, and the desktop AppImage. Package names were verified by installing them in the real base image rather than guessed: `cudnn9-cuda-12-8` is unusable there (it pins `libcudnn9-static` at 9.8.0.87 against the repo's 9.25.1.1), so the recipes use `libcudnn9-dev-cuda-12` + `libcudnn9-cuda-12`. The AppImage needed more than a flag: linuxdeploy follows only `DT_NEEDED`, which reaches cuDNN's 128 KB dispatcher but not the engine libraries it dlopens, so cuDNN would have failed to initialise on users' machines and quietly fallen back. The existing linuxdeploy shim gained a force-include directory; with the variable unset its argument list is byte-identical to before, verified by compiling the shim with `-Wall -Wextra` and testing the injection. `libcudnn_adv` is excluded — RNN/multi-head attention, which mold never asks cuDNN for, and it would add 259 MB. The `cargo install` line in the release notes deliberately does **not** get the feature: it runs on a user's machine, where the cuDNN SDK may not exist. ## Also fixed A pre-existing candle correctness bug the work unmasked: the im2col non-contiguous-kernel branch matmul'd the original strided kernel instead of the packed copy. `conv2d_grad_noncontiguous_kernel_gpu` fails on the old fork rev with `--features cuda` and passed only because cuDNN masked it. 🤖 Generated with [Claude Code](https://claude.ai/code) https://claude.ai/code/session_01AjPWQj3ZabxXoMRhptD4c9

jamesbrink added 7 commits August 28, 2026 08:39

chore: ignore .claude/ local agent workspace

e74f25f

style: rustfmt neighborhood_attention

613077a

Carried in unformatted from the NA3D commits this PR brings to main; the fork's Rustfmt CI job fails on it. Claude-Session: https://claude.ai/code/session_01AjPWQj3ZabxXoMRhptD4c9

jamesbrink merged commit 744ae3b into main Aug 31, 2026
11 checks passed

jamesbrink deleted the feat/cudnn-conv-dispatch branch August 31, 2026 05:00

jamesbrink mentioned this pull request Aug 31, 2026

perf(wan): render video convolutions on cuDNN, 2.1x faster VAE decode utensils/mold#1485

Merged

Navigation Menu

Repository navigation

perf(cuda): make the cuDNN convolution path usable, and 4.4x faster - #11

perf(cuda): make the cuDNN convolution path usable, and 4.4x faster#11
jamesbrink merged 7 commits into
mainfrom
feat/cudnn-conv-dispatch

jamesbrink commented Aug 31, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

jamesbrink commented Aug 31, 2026

Why this is 5 commits, not 1

The cuDNN work

Policy

A pre-existing correctness fix

Tests

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant