feat(z-image): run Z-Image attention on SageAttention when enabled

[Pfannkuchensack]

Note

Stacked PR 3 of 8: merge after #9659. 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, 6c8718945a.

Summary

With attention_backend: sage, Z-Image denoising enters the SageAttention scope from #9657, so its eligible attention calls run on SageAttention. Z-Image Control passes an attention mask and therefore stays on PyTorch SDPA, like everything outside the loop. The default auto is unchanged. The docs page lists Z-Image under Supported models.

Related Issues / Discussions

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

QA Instructions

Z-Image FP8, model loaded, three arms per seed as described in #9657:

	1024²	2048²
Step, default SDPA	520 ms	3184 ms
Step, cuDNN-first SDPA	490 ms	2819 ms
Step, SageAttention	462 ms (−11 %)	2231 ms (−30 %)
Whole generation	10.4 → 10.0 s (−4 %)	33.7 → 25.9 s (−23 %)
Peak VRAM in the denoise scope	unchanged	unchanged
LPIPS against default SDPA, median (cuDNN-first)	0.154 (0.092)	0.168 (0.112)

• Z-Image reacts strongly to any kernel change: even the exact cuDNN swap moves LPIPS to 0.09–0.11, and SageAttention stays within 1.7× of that. Blind side-by-side review: 8 of 8 pairs equal; the pair with the largest LPIPS (0.25) shows the same scene, sharpness and detail.
• Per step, the 32 calls of the main and noise-refiner layers run on SageAttention; the two text-only context-refiner layers are below the size threshold and stay on SDPA.
• Accuracy on real activations, every call of step 0 (32 calls): worst cosine 0.9970, relative L2 0.043.
• The timings, LPIPS and blind review come from the measurement prototype, which still used sageattn()'s default K smoothing; the selection feat(attention): add opt-in SageAttention backend for diffusion models #9657 ships is faster per call.

To try it: install SageAttention as the docs page from #9657 describes, set attention_backend: sage, restart and generate with Z-Image. 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

No material findings for this call site. Z-Image Control gains nothing until its mask is handled (possible follow-up).

Compatibility / Rollout

Merge after #9659. 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)

[lstein]

LGTM