feat(krea2): run Krea-2 attention on SageAttention when enabled

[Pfannkuchensack]

Note

Stacked PR 5 of 8: merge after #9661. It targets main, so until the PRs before it merge, its diff also shows their commits and the merges that carry them forward. This PR's own change is one commit, 482e7f348f.

Summary

With attention_backend: sage, Krea-2 denoising enters the SageAttention scope from #9657, so its eligible attention calls run on SageAttention; its grouped-query attention (48 query and 12 K/V heads) runs natively, without expanding K/V. While INVOKE_KREA2_SDPA_BACKEND pins a PyTorch kernel for measurement, the scope stays closed so the pinned kernel is what runs. The default auto is unchanged. The docs page lists Krea-2 under Supported models.

Related Issues / Discussions

Stacked on #9657 (SageAttention backend) and #9661, the PR before it.

QA Instructions

Krea-2 Turbo GGUF Q4_K_M, model loaded, arms as described in #9657:

	1024²	2048²
Step, SDPA (Krea-2's own cuDNN-first policy)	1254 ms	4976 ms
Step, SageAttention	1214 ms (−3 %)	4186 ms (−16 %)
Whole generation	11.5 → 11.0 s (−4 %)	42.7 → 36.2 s (−15 %)
Peak VRAM in the denoise scope	unchanged	unchanged
LPIPS against SDPA, median (no floor, see below)	0.058	0.238

• Krea-2 already ranks cuDNN first in its own SDPA policy, so the cuDNN-first arm is identical to the default and gives no noise floor: every composition change shows up in LPIPS. The pairs with the largest LPIPS (0.305 at 2048², 0.192 at 1024²) show the same scene with a shifted crop or a missing stack of books; hands and fabric are equally sharp at 1:1 and the lettering is correct.
• Blind side-by-side review: 8 of 8 pairs equal.
• Accuracy on real activations: block 0 has the same Q/K channel outlier as Qwen-Image (cosine 0.53 with K smoothing, 0.988 without, as feat(attention): add opt-in SageAttention backend for diffusion models #9657 ships).
• Per step, 28 of the 32 attention calls run on SageAttention; the other 4 are below the size threshold.
• The timings, LPIPS and blind review come from the measurement prototype, which still had K smoothing on; with the shipped selection one generation measured LPIPS 0.055 against SDPA.

To try it: install SageAttention as the docs page from #9657 describes, set attention_backend: sage, restart and generate with Krea-2. The first generation logs SageAttention ... serves diffusion-model attention on cuda:0 (...) with the GPU and kernel name, and with log_level: debug every generation logs how many calls ran on SageAttention and why the others did not. With attention_backend: auto the output is bit-identical to main.

Checks: ruff clean; pytest tests/app/invocations passes with the whole stack applied (2090 passed). No call-site test was added: this change only enters the scope, whose behavior #9657 tests.

Review

Resolved: with INVOKE_KREA2_SDPA_BACKEND set, SageAttention served the calls the override was meant to measure; the scope now stays closed while it is set.

Compatibility / Rollout

Merge after #9661. No API or persisted-state change; only attention_backend: sage behaves differently.

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)

[JPPhoto]

Looks good.