perf(krea2): rank cuDNN first for SDPA, with a fallback list and a probe - #172
Draft
Pfannkuchensack wants to merge 4 commits into
Draft
perf(krea2): rank cuDNN first for SDPA, with a fallback list and a probe#172Pfannkuchensack wants to merge 4 commits into
Pfannkuchensack wants to merge 4 commits into
Conversation
cuDNN is ~1.6x the memory-efficient kernel on the Krea-2 attention shape at identical peak memory. Measured end to end, that is 4-8% per generation, growing with resolution. set_priority=True is the mechanism, not a detail: without it sdpa_kernel() only *permits* the listed backends and torch still picks by its own order, in which cuDNN ranks last. The ranking would be a silent no-op. This stays a ranked list with fallbacks rather than a named default. An unavailable backend is skipped by the dispatcher, so the list degrades on its own -- to `efficient` where cuDNN is unusable, which is exactly today's behaviour, and on ROCm, where cuDNN is absent and flash rejects the additive mask. FLASH is deliberately kept in the list. Dropping it where a CUDA probe shows it absent would be actively wrong on ROCm, where flash is present and cuDNN is not: a dead entry in a ranked list costs nothing, a missing one costs a platform. Also here: - INVOKE_KREA2_SDPA_BACKEND, an opt-in override. Unset -- the only state a user sees by default -- is the ranked list, unchanged. The exclusive modes are the point: a run that completes proves that kernel was used, because an unavailable backend raises visibly instead of degrading to math unnoticed. Resolved once per generation and handed to every processor, not re-read per attention call. - A per-step benchmark in the denoise loop, built only when that variable is set and the device is CUDA. Its cuda.synchronize() pairs serialise the loop, so the default path must not reach them at all. The peak counter is reset immediately before the loop, so the number reported is the loop's and not the model load's; the first step is reported separately because it absorbs kernel selection and allocator warmup. - A startup probe of SDPA backend availability, logged once. Availability is not a property of the torch version: it depends on the build, the device and the shape, so the probe uses the shape the transformers actually attend over and the log line names it. Any failure is swallowed -- a diagnostic must never keep the server from starting. Not verified on sm_86. An RTX 3060 under torch 2.7.1 had cuDNN raise on this shape, while the same card under torch 2.11 did not; whether that is about sm_86, the shape or a version floor is not established. The fallback makes it a non-issue in principle, and one priority-cudnn generation on such a card should confirm it before this merges.
The ranked list is only safe because an unusable backend is skipped rather than raised on. That is a dispatcher property, and nothing in the suite held it: if a future torch made a ranked list raise instead, every device where cuDNN cannot serve would go from "today's behaviour" to a failed generation, silently as far as CI is concerned. fp32 is refused by the fused kernels, which makes cuDNN unusable on a card where it otherwise works -- a portable stand-in for the sm_86 report. Under it, an exclusive cuDNN selection raises `No available kernel` while the shipped ranked list completes. Verified on an RTX 4090; skipped without CUDA.
This was referenced Aug 28, 2026
The previous commit ranked cuDNN first, measured on a Windows 4090 where
flash is not compiled into the build at all. A probe of the same shape on a
Linux 30-series card shows the ordering is build-dependent, and that
cuDNN-first would have been a small regression there:
4090 / Windows 30-series / Linux
flash not compiled in 19.74 ms
cudnn 3.72 ms 21.27 ms
efficient 5.92 ms 31.45 ms
unranked (today) 5.89 ms 20.15 ms
Two things follow. Flash is the fastest kernel where the build has it, and
is already what runs today -- without set_priority torch picks by its own
order, in which flash outranks efficient, so the unranked call lands on
flash (20.15ms) rather than efficient. And cuDNN is where the win actually
comes from: flash refuses the additive padding mask the regional-prompting
blocks pass, so on those blocks it is skipped and cuDNN takes over at
1.6x-2.0x over efficient -- and on a build without flash, that is every
block.
Flash-first gives up nothing where flash cannot run. Re-measured on the
4090 with the shipped order: unmasked 5.252ms -> 3.247ms, masked 6.657ms ->
3.579ms, against 3.237ms for cuDNN-first. Identical within noise, because
flash is skipped there.
This also corrects a claim carried over from the design notes: that today's
list and an exclusive [EFFICIENT] are the same thing. That was measured on
the flashless build and does not hold on Linux CUDA.
Mutations verified as caught: restoring cuDNN-first (3 tests), moving
efficient above cuDNN (2), removing flash (5), moving math off the end (1),
dropping set_priority (2).
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Performance.
Krea2MemoryEfficientAttnProcessorruns SDPA under a permissive backend list and lets torch choose. Two of the four backends it could be using are much faster than what it settles on, and which one wins depends on the build. This makes the list ranked instead of merely permissive. Measured through the real stack on an RTX 4090: 1.079× at 1536², 238 ms per step, with peak VRAM identical to the byte.Per attention call, on the real Krea-2 shape (
[1, 48, 4608, 128], bf16):The two columns disagree about which kernel wins, and the shipped order encodes both:
Measured on the 4090 with flash absent, confirming flash-first gives nothing up where flash cannot run: unmasked 5.252 ms → 3.247 ms (1.62×), masked 6.657 ms → 3.579 ms (1.86×) — identical to what cuDNN-first produced.
set_priority=Trueis the mechanism, not a detail.sdpa_kernel([...])without it only permits the listed backends; PyTorch still picks by its own internal order, in which cuDNN ranks last and is never chosen. AddingCUDNN_ATTENTIONto the list withoutset_priority=Truewould be a silent no-op.It stays a ranked list with fallbacks, never a named default. An unavailable backend is skipped by the dispatcher, so the list degrades on its own:
flash— unchanged from today, and the fastest cellcudnncudnncudnnefficient— today's behaviourefficientflashefficientFLASHis deliberately kept — and leads. A sibling proposal was to drop it where a probe shows it absent. That would be wrong twice over: on ROCm flash is available and cuDNN is not, and on Linux CUDA it is the fastest kernel of the four. A dead entry in a ranked list costs nothing; a missing one costs a platform. A test pins this so it does not get "cleaned up" later.Also in this PR
INVOKE_KREA2_SDPA_BACKEND, an opt-in override. Unset — the only state a user sees by default — is the ranked list, unchanged.cudnn/efficient/flash/mathpriority-cudnnThe exclusive modes are the point: a run that completes proves that kernel was actually used, because an unavailable backend raises visibly instead of quietly degrading to
math. Invalid values raise and name the valid ones. Resolved once per generation inbuild_krea2_attention_processorsand handed to every processor — not re-read per attention call, of which there are dozens per step.A per-step benchmark in the denoise loop, built only when that variable is set and the device is CUDA. That gate is not cosmetic: the
cuda.synchronize()pair around each step serialises the loop, so the default path must not reach it at all. The peak counter is reset immediately before the loop, so the number reported is the loop's and not the model load's, and the first step is reported separately because it absorbs kernel selection and allocator warmup.A startup probe of SDPA backend availability, logged once next to the existing
check_cudnn. Availability is not a property of the torch version: ROCm builds have no cuDNN attention, Windows CUDA builds usually have no flash, and both depend on the device architecture and the shape. So the probe askscan_use_*against the shape the transformers actually attend over — 24 heads, seq 1024, head_dim 128 — and the log line names that shape, because a line without it would be unactionable. It also saysno mask: flash refuses the additive mask the regional-prompting paths pass, so this is a diagnostic, not a dispatch table. It is a capability query, not a benchmark — no kernel launches, no timing, and the probe tensor is released afterwards. Every exception is swallowed: a diagnostic must never be the reason the server does not start.Related Issues / Discussions
The design notes live in the local
.ideas/krea2-sdpa-cudnn-fallback.mdand.ideas/attention-backend-selection.md. This PR is the concrete version of the latter's PR 1 and PR 2; its PR 3 (moving the two hardcoded backend lists into anAttentionFacet) depends on the architecture registry and is not part of this.QA Instructions
All of the following was measured on this branch, on an RTX 4090 (sm_89), torch 2.7.1+cu128, driver 610.47, Windows, against a real InvokeAI server and the real model manager.
End to end, through the real stack
Krea-2 Turbo GGUF Q4, 8 steps, CFG 1.0, fixed seed, no VAE decode in the graph, 1 warmup + 3 measured generations per mode, one server per mode. Transformer 100 % VRAM-resident in every cell (
VRAM: … (100.0%)), so no partial loading distorts the comparison.1536×1536 — the decisive cell:
efficient(= today)priority-cudnn(= ships)The effect is 238 ms per step against a within-mode spread of 13–36 ms — roughly seven times the noise. Peak VRAM is identical to the byte, so this is pure compute, not a memory trade.
1024×1024 — a null result, recorded deliberately:
efficientpriority-cudnncudnn(exclusive)This cell shows cuDNN slower, and it should not be read as evidence against the change. The claimed effect at 1024² is ~57 ms while the individual runs spread over 190 ms (1331–1520), so the measurement cannot resolve it in either direction. It was produced by running three modes back-to-back in one sitting — the design the design notes explicitly warn about at this effect size, having hit the same null in the same cell. The 1536² cell, where the effect is 4× larger and the spread 5× smaller, is what carries the conclusion. Both are reported here rather than only the favourable one.
Isolated kernel, the shape the processor actually runs
bf16, 48 heads (K/V already expanded by
repeat_interleave), head_dim 128, 5 warmups + 30 timed, old list against new:Peak allocation is identical in every cell. Outputs agree at corr ≥ 0.9999979, maxdiff ~1e-3 — bf16 rounding. Worth noting because the design notes only measured the unmasked case: the masked path, which regional prompting takes, benefits most at low sequence length (1.98×).
Those figures were taken with cuDNN ranked first, on a build without flash. Re-measured with the shipped flash-first order on the same machine, the numbers are unchanged (3.247 ms against 3.237 ms unmasked) — flash is skipped there, so the two orders resolve identically.
The fallback, verified rather than assumed
The whole design rests on one dispatcher property. Reproduced by making cuDNN unusable for a call on a card where it otherwise works:
can_use_cudnn_attentionset_priority=TrueRuntimeError: No available kernelRuntimeError: No available kernelTwo CUDA-gated tests pin this (
TestTheFallbackIsReal). They are skipped on CPU-only CI but do run and pass on a CUDA machine — verified, not merely written.Startup probe, on real hardware
flash=nois correct: this Windows CUDA build was not compiled with flash attention. On CPU the probe returnsNoneand logs nothing.Test suite
tests/backend/krea2/test_sdpa_backend_selection.py— 26 tests (24 portable + the 2 CUDA-gated fallback tests).tests/app/util/test_attention_backend_probe.py— 7 tests.set_priority(2 failures), restoring cuDNN-first (3), moving efficient above cuDNN (2), removingFLASH(5), movingmathoff the end (1).tests/backend/krea2 tests/app/utilon CUDA — 209 passed, 0 skipped on the 4090. The two files touched by this PR were also run on an sm_86 card: 29 passed, 0 skipped.ruff check/ruff format --checkclean;pytest --collect-onlyclean; noopenapi.jsonchange.Merge Plan
On the sm_86 concern. An RTX 3060 (sm_86) under torch 2.7.1 was reported to have cuDNN raise
RuntimeErroron this shape, while a 4090 (sm_89) under the same torch does not. That was the reason this PR was opened as a draft. Two independent pieces of evidence have since narrowed it:On a 4090, by making cuDNN unusable:
sdpa_kernel([backend])— exclusive, single-element, no fallback. That mode raisesRuntimeError: No available kernelwhenever the backend is unusable, which is reproduced above. So the report is evidence about cuDNN availability on that card, not about a ranked list failing.On an actual 30-series (sm_86) card, Linux, Python 3.12: a shape probe at the real Krea-2 shape shows cuDNN working there — 21.27 ms, no
RuntimeError— which is the case the original report warned about, and it does not reproduce. That probe is also what surfaced flash leading on that card, and is why the order above is flash-first rather than cuDNN-first. Separately, the full test suite runs green there — 29 passed, including bothTestTheFallbackIsRealtests. That matters because those two do not skip silently: the helper callspytest.skipif cuDNN can serve the probe call, so their having run is itself proof that cuDNN was genuinely unavailable there, and that the ranked list completed anyway while the exclusive selection raised. The dispatcher fallback is therefore verified on sm_86 itself, not inferred from sm_89.Where that leaves the concern. The failure the report described — cuDNN raising on this shape on sm_86 — does not reproduce: it completes in 21.27 ms there. And the fallback that would have covered it if it did is verified on the same architecture. What has not been run is a full Krea-2 generation on such a card, which is the literal wording of the original pre-merge check; the shape probe answers the same question without needing a model. The residual risk is now narrow enough that I would take this out of draft, but that is the reviewer's call.
Test coverage note
The tests were updated with the order. Five mutations were verified as caught: restoring cuDNN-first (3 failures), moving efficient above cuDNN (2), removing flash from the list (5), moving math off the end (1), and dropping
set_priority(2).One product question for the reviewer:
INVOKE_KREA2_SDPA_BACKENDis a user-visible environment variable whose exclusive modes fail hard on purpose. It ships here undocumented and off by default. If it should not be a supported surface, it can be dropped — but it is also what makes the A/B above reproducible by anyone, and what the sm_86 check would use.No DB schema, no redux slice, no dependency change, no API change. Independent of #171, which is in flight in parallel — the two touch disjoint files and can merge in either order.
Checklist
INVOKE_KREA2_SDPA_BACKENDis deliberately undocumented; see the question in Merge PlanWhat's Newcopy (if doing a release after this PR) — worth a line: Krea-2 generation is ~8 % faster at 1536² on CUDA cards with cuDNN attention🤖 Generated with Claude Code