[Pfannkuchensack]

Summary

PyTorch's own SDPA order is flash, memory-efficient, math, cuDNN. cuDNN ranks behind the always-eligible math kernel, so it is practically never chosen, and Windows CUDA builds have no flash kernel: every model ran on memory-efficient attention. Only Krea-2 and Anima opened their own sdpa_kernel windows to reach cuDNN, and those windows are process-global, so they could not hold while two GPUs run sessions at once.

This PR installs one process-wide order at startup for every model: flash, cuDNN, memory-efficient, math (attention_backend: auto).

• Windows CUDA: denoising is 4–18 % faster, peak VRAM unchanged. Images at the same seed change slightly (bf16 accumulation order), as with a PyTorch upgrade.
• Linux CUDA: flash stays first; only calls flash refuses (masks, wide heads, grouped-query) move from memory-efficient to cuDNN.
• ROCm: has no cuDNN attention, so nothing changes.
• Text encoders with grouped-query attention (Qwen-Image, Qwen-Image 2.1, ERNIE-Image, Ideogram 4, Anima) ran on the math kernel under PyTorch's order; they now run on cuDNN.
• cuDNN runtime failure: a non-OOM error inside cuDNN demotes cuDNN behind math for the rest of the process, logs one warning and retries the call (grouped-query calls are retried with expanded K/V, so the retry never lands on math).
• Krea-2 asks the dispatcher whether a fused kernel takes its unexpanded grouped-query K/V, so the answer follows both the enable flags and the order.
• Diagnostics: the SDPA availability line is logged once per GPU for every model (from the session processor, no longer from two denoise nodes) and names the active order, e.g. cudnn=yes flash=no efficient=yes math=yes; priority FLASH_ATTENTION > CUDNN_ATTENTION > EFFICIENT_ATTENTION > MATH > OVERRIDEABLE.
• Removed: the Anima and Krea-2 default windows and the now-unreachable foreign-window bookkeeping in sdpa_scope.py. INVOKE_KREA2_SDPA_BACKEND keeps working for benchmarking.
• The SageAttention page now states its gains against the new default.

Related Issues / Discussions

Related: #9529 (attention optimizations), #9831 (Wan 2.2 A14B at 720p).

QA Instructions

Tests

• CPU: CUDA_VISIBLE_DEVICES=-1 uv run --no-sync pytest tests/backend/util tests/backend/krea2 tests/app/util tests/app/services/session_processor → 1188 passed, 21 skipped.
• RTX 4090, Windows, torch 2.13.0+cu130, same suites: 1203 passed, 6 skipped (MySQL- or Linux-only). The CUDA-gated tests ran, including the grouped-query/order cases and kernel agreement against math.
• RX 9060 XT, Windows, torch 2.13.0+rocm10: SDPA and Krea-2 suites 171 passed, 3 skipped (cuDNN and fused grouped-query attention do not exist there).
• ruff check/format, invokeai/frontend/api contract tests, docs build, check-docs-data: pass.

Denoising time per generation, PyTorch's order → this order (RTX 4090, Windows; median of 12 generations at 1024, 4 at 2048 and video; Wan A14B interleaved A B B A A B, 3 runs each):

Model	1024×1024	2048×2048 / video
FLUX.1	14.8 → 13.8 s (−7 %)	78.3 → 64.8 s (−17 %)
FLUX.2 Klein 9B	2.8 → 2.6 s (−7 %)	13.8 → 11.8 s (−14 %)
Z-Image	4.3 → 4.0 s (−7 %)	26.7 → 23.1 s (−13 %)
Qwen-Image	55.1 → 52.9 s (−4 %)	184.7 → 159.1 s (−14 %)
SDXL	3.2 → 3.2 s (±0)	15.8 → 14.5 s (−8 %)
LTX-2 video, 1248	–	41.5 → 37.2 s (−10 %)
Wan 2.2 TI2V-5B, 1280×704×121	–	349.8 → 286.1 s (−18 %)
Wan 2.2 T2V A14B, 832×480×81, 4-step	–	86.1 → 72.2 s (−16 %)

Validation on this branch (RTX 4090, graphs built by the UI's own graph builder, one server process, order switched per run, warm-up then torch / new / new / torch):

• 10 families at 1024² (SDXL, FLUX.1, FLUX.2 Klein, Z-Image, Qwen-Image, Qwen-Image 2.1, Krea-2, Anima, ERNIE-Image, Ideogram 4): all complete, repeats bit-identical within each order, no NaN or blank images, images equivalent between orders (same composition, detail shifts). The kernel each call ran on was counted through the dispatcher and matches the order.
• Same order on this branch and on its base: bit-identical images for SDXL, FLUX.1, Z-Image, Krea-2 and Anima, so removing the windows and adding the fallback changes nothing else.
• Partial loading with device_working_mem_gb: 5, new shapes for cuDNN (the documented Windows cuDNN lazy-load hang): FLUX.1 int8 1536²/2048² 34.4/67.8 s, Z-Image int8 1536²/2048² 12.7/25.4 s, Qwen-Image 2.1 1536² 35.7 s; no stalls.
• RX 9060 XT: PyTorch's order vs this order, and base vs branch, bit-identical for SDXL and Krea-2.

Linux CUDA: tested on 2× RTX PRO 6000 Blackwell (Ubuntu 26.04, torch 2.13.0+cu130); results in #9955 (comment).

Not covered: other NVIDIA generations (Ampere), MPS, XPU. The cuDNN fallback's failure path is unit-tested with a failing stand-in kernel; a real cuDNN failure could not be provoked. Z-Image Turbo GGUF could not be run: it fails on main already (TypeError in aten.where on a GGMLTensor in diffusers' transformer_z_image._prepare_sequence), unrelated to this PR.

Review

Resolved:

• Krea-2's grouped-query check combined enable flags with can_use_*, which ignores the priority order; it now asks the dispatcher, so a "yes" cannot route an unexpanded call to math (~9 GB per call at 1024²) when the order is not in force.
• The dispatcher query swapped Python warning filters on the per-call path. PyTorch's warnings there are C++ stderr lines the filter does not catch, and each swap resets process-wide warning state; the filter is gone, and with math disabled (an exclusive override) the query uses the silent can_use_* checks.
• The fallback retried a failed grouped-query call unchanged, which would have landed on math; the retry expands K/V. Sticky CUDA errors (torch.AcceleratorError) and out-of-memory errors are not treated as cuDNN failures.
• A failed availability probe is remembered instead of re-allocating on every item; the probe runs after the cancel/skip checks and cannot fail an item.

Remaining risks:

• Every Windows CUDA attention call now goes through cuDNN. A GPU/driver/cuDNN combination that fails at run time is caught by the fallback; a hang during cuDNN's lazy kernel loading cannot be caught (none seen in the partial-loading sweep).
• The vendored PiD network binds scaled_dot_product_attention at import, before startup patching, so its decode attention bypasses the fallback (it already bypassed the ROCm guard).

Compatibility / Rollout

• On CUDA, images at the same seed differ slightly from earlier versions; worth a line in the release notes.
• attention_backend keeps its values (auto, sage); only its description changed. openapi.json, schema.ts and the generated settings docs are regenerated. No persisted data or migrations involved.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

[Pfannkuchensack]

Linux CUDA test run. Machine: Verda 2× NVIDIA RTX PRO 6000 Blackwell Server Edition (96 GB, sm_120), driver 580.178.04, Ubuntu 26.04.1, torch 2.13.0+cu130, cuDNN 9.20, Triton 3.7.1. The PR head was compared with its merge base on main; each GPU ran one lane, pinned by UUID.

• pytest (tests/backend/util tests/backend/krea2 tests/app/util tests/app/services/session_processor tests/test_config.py):
  • GPU visible: 1265 passed, 4 skipped.
  • -m slow: 24 passed.
  • CPU-only (CUDA_VISIBLE_DEVICES=-1): 1265 passed.
• E2E against the merge base (64acb44), through the API.
  • Cases: SDXL 1024²/25 steps, FLUX.1 dev int8 1024²/20, Z-Image Turbo 1024² and 2048²/8, Krea-2 Turbo int8 1024²/8.
  • Each case: 1 warm-up and 3 measured runs, in two interleaved rounds.
  • SDPA log: one availability line per server, cudnn=yes flash=yes efficient=yes math=yes, priority FLASH > CUDNN > EFFICIENT > MATH. No cuDNN demotion.
  • SDXL and FLUX are bit-identical to main.
  • Against main, Z-Image is at 29.9 dB (1024²) and 33.2 dB (2048²) PSNR, and Krea-2 at 26.5 dB. These are the calls flash refuses, which now go to cuDNN instead of the memory-efficient kernel. Composition and quality are unchanged; only fine detail moves.
  • Time against main, from a quiet re-run of 3 interleaved rounds:
    • SDXL 2.403 s vs 2.400 s, FLUX 7.699 s vs 7.687 s.
    • Krea-2 and Z-Image −0.3 %.
  • Repeats are bit-identical within each arm.

[lstein]

Adversarial review of b0f6bb4352: no blockers

I tried to break the cuDNN fallback's lock and demotion counter, the Krea-2 grouped-query check, the probe's move into the session processor, test isolation, and the ROCm and XPU paths. The core design holds:

• Exclusive windows can't undo a demotion. sdpa_policy uses sdpa_kernel without set_priority, which saves and restores only the enable flags, never the priority order.
• Krea-2 is safe without startup patching. If the order was never installed (tests, scripts), _serves_grouped_query_attention gets "math" from the dispatcher and expands K/V, so an unexpanded grouped-query call never lands on math.
• The probe can't stall a peer GPU. Its empty_cache() goes through the peer-aware wrapper.
• Two workers failing at once are handled. The lock plus _cudnn_demotions means exactly one demotes and warns, and the other retries. Lock order with _order_lock is consistent.
• Test isolation is fine. _cuda_build and preserve_sdpa_priority_order restore the global state they change.

Findings

1. PiD still bypasses the fallback, and the fix is one line. invokeai/backend/pid/_src/networks/pixeldit_official.py:272 calls the bare scaled_dot_product_attention imported at line 18, so it gets the unpatched function. Line 707 of the same file already calls F.scaled_dot_product_attention. Changing line 272 to match puts PiD decode attention behind both the cuDNN fallback and the ROCm wide-head guard.

2. One failure demotes cuDNN for every model (design trade-off). A single cuDNN RuntimeError on any one shape, for example in one text encoder, ranks cuDNN behind math for the rest of the process. That gives up the whole 4–18 % gain for every model, with one warning as the only sign. Fine for now; a follow-up could scope the demotion to the failing call's shape.

3. The warning can be misleading. Some execution failures arrive as a plain RuntimeError but have already poisoned the context, for example CUDNN_STATUS_EXECUTION_FAILED (as opposed to torch.AcceleratorError). Those demote cuDNN and log "generation continues", and then the retry fails on the poisoned context. Low impact.

4. Rare race, benchmarking only. It needs INVOKE_KREA2_SDPA_BACKEND set and two GPUs running sessions:

   • Session A opens an exclusive efficient window.
   • Session B's call runs on memory-efficient (the only enabled kernel) and fails for an unrelated reason.
   • A's window closes before B calls sdpa_backend_choice.
   • The dispatcher now names cuDNN, so B demotes cuDNN by mistake.

   Normal users can't hit this.

Remaining risk

The fallback only catches errors. If cuDNN returns wrong values or NaN, nothing catches it, and Ampere and Turing weren't tested. Every Windows CUDA attention call now goes through cuDNN, so the release-note line about images changing at the same seed matters.

This was a static review: I did not run the PR's tests.

[Pfannkuchensack]

@lstein Thanks for the thorough pass.

1. PiD: fixed in 7c62258. Decode attention now calls F.scaled_dot_product_attention, so it goes through the cuDNN fallback and the ROCm guard. No other module in invokeai/backend binds SDPA at import time.
2. Process-wide demotion: agreed it's the conservative choice. I'll keep it for this PR and open a follow-up to scope the demotion to the failing call's shape if failures show up in the field.
3. Poisoned-context failures: right, an error like CUDNN_STATUS_EXECUTION_FAILED that arrives as a plain RuntimeError gets demoted and retried. If the context is already broken, the retry raises, and the generation fails with that error, so the log may say "generation continues" right before the failure. I left out message matching for specific cuDNN statuses since it would be brittle across torch versions.
4. Exclusive-window race: it needs INVOKE_KREA2_SDPA_BACKEND plus two GPUs, and the effect is only that cuDNN gets demoted too early. I left it as is, since the variable is a benchmarking aid.

On the remaining risk: agreed, the fallback only catches errors, not wrong results. The release note will say that images at the same seed change slightly.

[lstein]

Approving at f2893206ca.

The PiD fix in 7c62258 is correct: decode attention now goes through the patched F.scaled_dot_product_attention, and nothing else in invokeai/ binds SDPA at import time. Your answers to findings 2–4 are fine as accepted trade-offs, with the shape-scoped demotion as a follow-up.

I merged main into the branch (f289320) to clear the conflicts. All three were mechanical:

• sage-attention.mdx: kept this PR's new timings, plus main's updated Z-Image note ("Also serves ControlNets inside the denoising loop", from #9943).
• scripts/calibrate_flux2_working_memory.py: took main's rewrite (#9927) and kept the install_default_sdpa_priority() call and its import.
• openapi.json: took main's config description and re-applied this PR's attention_backend wording. Compared with main, the only difference in openapi.json/schema.ts is that wording.

Checks on the merge: scripts/generate_docs_json.py reproduces settings.json, ruff passes, and tests/backend/util tests/backend/krea2 tests/app/util tests/app/services/session_processor gives 1347 passed, 7 skipped (ROCm, torch 2.13.0+rocm10).

Nit, non-blocking: the install_cudnn_sdpa_fallback docstring (invokeai/backend/util/attention.py, the paragraph naming pixeldit_official.py) and the PR description's "Remaining risks" still say PiD bypasses the fallback.