[Pfannkuchensack]

Summary

Builds on #9955 (one process-wide SDPA order, cuDNN ranked ahead of memory-efficient attention). Until #9955 merges, this branch contains its commits; the change of this PR is the last commit, f8a93c0.

#9955 gives attention_backend: auto a fixed order (flash, cuDNN, memory-efficient, math). This PR lets the server admin pin one PyTorch kernel instead, in invokeai.yaml:

attention_backend: efficient   # auto | sage | cudnn | flash | efficient | math

• The pinned kernel ranks first, followed by auto's order. It is never exclusive and never raises: a call it cannot serve runs on the next kernel.
• Every such call is reported once per node type, pinned kernel, serving kernel and reason, e.g.
  attention_backend 'efficient' is not possible for attention in qwen_image_2_1_text_encoder: the memory-efficient kernel cannot serve this call (bfloat16, head_dim 128, no mask, 32 query / 8 key-value heads, enable_gqa, causal); it runs on CUDNN_ATTENTION. PyTorch's own reason follows once on stderr. Other reasons: the build or GPU has no such kernel (e.g. flash on Windows CUDA, cudnn on ROCm, also warned at startup), cuDNN was demoted after a runtime failure, an exclusive INVOKE_KREA2_SDPA_BACKEND window switched it off, or ROCm's wide-head guard sent the call to math.
• The node type comes from a context variable the session processor sets around each node, so every model family is covered without per-family code.
• Krea-2 expands its grouped-query K/V when that lets the pinned kernel serve the call, so a pinned efficient really runs Krea-2 on memory-efficient attention instead of cuDNN.
• math materializes the full score matrix (slow, much more memory) and is meant for debugging; the working-memory estimate follows the dispatcher and prices it.
• auto and sage install nothing new and behave exactly as with perf(attention): rank cuDNN ahead of memory-efficient attention for every model #9955. With sage, the calls SageAttention does not serve (text encoders, VAEs, masked and short calls) use auto's order.

Use case: an escape hatch from cuDNN (e.g. a GPU/driver combination where it misbehaves or hangs, which #9955's automatic fallback cannot catch), and benchmarking or support questions.

Related Issues / Discussions

Depends on #9955.

QA Instructions

Tests

• CPU, torch 2.13: CUDA_VISIBLE_DEVICES=-1 uv run --no-sync pytest tests/backend/util tests/backend/krea2 tests/app/util tests/app/services/session_processor tests/test_config.py → 1288 passed, 35 skipped.
• CPU, torch 2.7.1 (the macOS / Linux aarch64 pin; its own SDPA order differs): same SDPA, Krea-2, probe, caller and config suites → 637 passed, 27 skipped.
• RTX 4090, Windows, torch 2.13.0+cu130: 1315 passed, 8 skipped (Apple Silicon, ROCm-only, MySQL- or Linux-only), including the CUDA-gated report, demotion, grouped-query and startup-context tests.
• RX 9060 XT, Windows, torch 2.13.0+rocm10 (AOTriton experimental kernels on): SDPA and Krea-2 suites 232 passed, 7 skipped (platform gates), including the ROCm wide-head reason.
• ruff check/format, invokeai/frontend/api contract tests: pass. openapi.json, schema.ts and the generated settings docs are regenerated.

End to end (graphs built by the UI's own graph builder; the kernel each call ran on counted through the dispatcher):

Setting	GPU	Result
efficient	RTX 4090	SDXL, FLUX.1, Z-Image, Qwen-Image 2.1 and Krea-2 (expanded heads) run on memory-efficient attention; Qwen-Image 2.1's grouped-query text encoder runs on cuDNN and is reported once
math	RTX 4090	SDXL 768² runs entirely on math, nothing reported
flash	RTX 4090	startup warning (no flash in Windows builds); each node type reported once, runs on cuDNN (VAEs on memory-efficient)
cudnn	RX 9060 XT	startup warning (ROCm has no cuDNN attention); text encoder and denoise run on flash, the VAE's 512-wide heads on math (ROCm guard), each reported once
auto	RTX 4090	bit-identical to #9955 for SDXL, FLUX.1, Z-Image, Krea-2, Anima; no report lines
auto	RX 9060 XT	bit-identical to #9955 for SDXL and Krea-2
sage	RTX 4090	SageAttention serves FLUX.1 1140/1140, Krea-2 1120/1280, Z-Image 256/272, SDXL 250/3500 calls; the rest runs on cuDNN

Repeats were bit-identical within every setting.

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

Not covered: other NVIDIA generations (Ampere), macOS (the report skips devices the dispatcher cannot answer for; the MPS test needs Apple Silicon), XPU.

Review

Resolved:

• Krea-2 passed grouped-query K/V unexpanded whenever any fused kernel took them, so a pinned efficient still ran Krea-2 on cuDNN; it now expands when the expanded call gets a higher-ranked kernel.
• That rule kept K/V grouped when both variants land on math. On ROCm, math runs in chunks there and PyTorch re-expands K/V for every chunk (~1.8 GiB of copies per call at 1024² instead of 108 MiB once); grouped heads now require a fused kernel, as before.
• auto did extra dispatcher work per Krea-2 call; with flash ranked first (auto, flash, PyTorch's order, after a cuDNN demotion) the decision is the previous one, unchanged.
• Report reasons for a demoted cuDNN and for the ROCm wide-head guard gained tests; the config Literal and the kernel mapping are pinned against each other.

Remaining limits:

• The vendored PiD network binds scaled_dot_product_attention at import, before startup patching: its attention follows the pinned order but is not reported.
• A kernel-availability probe could cache a wrong "no kernel" if another session's exclusive INVOKE_KREA2_SDPA_BACKEND window toggled twice within the probe; that variable is a benchmarking tool.

Compatibility / Rollout

• attention_backend gains four values; the default auto is unchanged, and the setting stays startup-only (not runtime-writable). The OpenAPI schema and generated types change accordingly; the frontend does not read this field.
• 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 (same suites as perf(attention): rank cuDNN ahead of memory-efficient attention for every model #9955):
  • GPU visible: 1324 passed, 8 skipped.
  • -m slow: 24 passed.
  • CPU-only: 1324 passed.
• auto: bit-identical to perf(attention): rank cuDNN ahead of memory-efficient attention for every model #9955 on all 5 cases. Time is within 2 % of main.
• Pinned cudnn, flash, efficient (and math for SDXL/Z-Image):
  • Every run completes.
  • There is no startup warning for flash or cudnn on this build.
  • Each "not possible" report appears once per (node type, kernel, reason). Example: cuDNN cannot serve the VAE attention at head_dim 384/512, so those calls run on EFFICIENT_ATTENTION.

[lstein]

Reviewed at ce27916ebc, this PR's own delta on top of #9955 (fc00e2bcde..ce27916ebc). The design holds up well overall. One change is requested, plus one nit.

1. attention_backend: math is not budgeted for most model families (change requested)

invokeai/backend/util/attention.py (sdpa_score_matrix_bytes)

The PR says "the working-memory estimate follows the dispatcher and prices it". That holds only where an estimator calls sdpa_score_matrix_bytes:

• FLUX.2, Krea-2, Qwen-Image, Qwen-Image-2.1 and Z-Image denoise
• the VAE paths (vae_working_memory.py), the Qwen-Image-2.1 text encoder and MOSS

SD1.5/SDXL, FLUX.1, Anima, Wan, LTX-2, CogView4, ERNIE, MiniMax H3 and Ideogram4 reserve nothing for a score matrix.

How to trigger it: on CUDA, set attention_backend: math and run FLUX.1 at 1024². That is 24 heads × 4608² × 17 B ≈ 8.1 GiB per attention call, all of it unreserved, so the model cache under-reserves and the generation OOMs. Wan 2.1 at 480p with 81 frames (~32.8k tokens, 40 heads) needs about 680 GiB per call. On ROCm the guard's chunking bounds this, so the problem is CUDA-only.

math is documented as a debugging mode that "needs much more memory", but the claim that it is priced reads as a guarantee. Any one of these would resolve it:

• Narrow the PR text and the attention_backend description to the families that are actually priced.
• Log a startup warning when math is pinned, saying the working-memory estimates don't cover it for the other families.
• Price it in the remaining estimators.

2. Nit: attention that runs on the CPU is reported

_pinned_kernel_bypass returns ("device", "cpu"). With cudnn or efficient pinned, a CPU-only install logs one warning per node type ("on cpu PyTorch picks the kernel itself, without the ranking"). So does a GPU install whenever a node's attention runs on the CPU. The message is accurate, but the pin was never meant to apply on the CPU. Skipping the report for CPU tensors, the same way MPS is already skipped, would remove the noise.

Checked and found sound

• Wrapper order (ROCm guard → report → cuDNN fallback → sage): guard chunks aren't reported once per chunk, and a retry after cuDNN is demoted is reported as demoted.
• cuDNN demotion while cuDNN is pinned: flash ends up first again, so Krea-2 goes back to its old decision.
• Krea-2 grouped-vs-expanded decision under each pin: on ROCm (W7900, torch 2.13+rocm10) I compared the expanded-call probe with the real call for efficient/cudnn/math × 3 sequence lengths (odd lengths included) × 4 mask shapes. They never disagreed.
• attention_caller: there is a single node-execution site, and the value is reset when a node raises.
• PyTorch's debug reason: it goes to stderr as a C++ warning, so the global SilenceWarnings in SD/SDXL denoise doesn't hide it.
• Calibration: attention_backend is already part of the calibration fingerprint.
• Tests: the focused suites pass on ROCm Linux (313 passed, 7 skipped).

Not checked here: CUDA-only paths (flash missing on Windows, cuDNN demotion, the probe on cuDNN) and XPU.

[Pfannkuchensack]

@lstein Thanks, both addressed in ae27189 (also merged with main).

1. math and working memory: the math kernel's score matrix is only reserved by the estimators that call sdpa_score_matrix_bytes: FLUX.2, Krea-2, Qwen-Image, Qwen-Image-2.1, Z-Image denoise, the VAE paths, the Qwen-Image-2.1 text encoder and MOSS. Pinning math now logs a startup warning. It calls math a debugging mode and says most families' estimates leave the score matrix out, so generations may run out of memory. On ROCm, the SDPA guard still splits large non-causal math calls into chunks. The config description and the Attention Kernel docs say the same, and no longer suggest that math is reserved.
2. CPU reports: the report now skips CPU tensors, as it already did for MPS. A CPU-only install, or CPU attention on a GPU install, no longer logs a warning per node type. The tests cover both the skip and the report logic.

[lstein]

Re-reviewed at 4b951225d1. Both points from the previous review are resolved in ae27189c42, and the other new commits are merges that add nothing else to this PR's changes. Approving.

1. math and working memory. Pinning math now logs a startup warning. The warning, the attention_backend description, the generated docs and the OpenAPI schema all say that most families' working-memory estimates leave out the score matrix. The warning is only logged once the pinned order is actually in place, which is the right condition.
2. CPU reports. The report now covers only cuda/xpu. CPU attention is skipped, so it is no longer reported per node type and no longer pays for the per-call kernel query. The replacement test covers the skip, and the remaining report tests treat the CPU as a stand-in GPU.

Checked:

• The focused suites (pinned report, sdpa_scope, Krea-2, node attribution, config, cuDNN fallback, ROCm guard, backend choice) pass on ROCm Linux (W7900, torch 2.13+rocm10): 317 passed, 7 skipped.
• The new tests catch the bugs they cover: putting the CPU back on the report list, or removing the math warning, makes the matching new test fail.

Nit (description only, no code change needed): the PR body still says the math working-memory estimate "follows the dispatcher and prices it". Please reword that to match the new warning.

Not checked here: the CUDA-only paths and XPU.