Skip to content

feat(sdxl): run SDXL attention on SageAttention when enabled - #9663

Draft
Pfannkuchensack wants to merge 24 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/sage-attention-sdxl
Draft

Pfannkuchensack wants to merge 24 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/sage-attention-sdxl

Conversation

@Pfannkuchensack

@Pfannkuchensack Pfannkuchensack commented Oct 3, 2026 •

Copy link
Copy Markdown
Member

Note

Stacked PR 6 of 8: merge after #9662. 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, 1968eb9c58.

Summary

With attention_backend: sage, SDXL denoising and SDXL tiled multi-diffusion enter the SageAttention scope from #9657, so eligible attention calls run on SageAttention, including ControlNets that run inside the loop. SD1.5 and SD2 do not enter it: SD1.5's head sizes (40, 80, 160) never qualify, and SD2 is unmeasured. With attention_type: sliced, SDXL's attention does not go through SDPA, so nothing changes there. The default auto is unchanged. The docs page lists SDXL under Supported models and notes the sliced case.

Related Issues / Discussions

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

QA Instructions

SDXL, model loaded, three arms per seed as described in #9657:

1024² 2048²
Step, default SDPA 134 ms 631 ms
Step, cuDNN-first SDPA 135 ms 579 ms
Step, SageAttention 126 ms (−6 %) 496 ms (−21 %)
Whole generation 4.4 → 4.2 s (−5 %) 19.6 → 16.1 s (−18 %)
Peak VRAM in the denoise scope unchanged unchanged
LPIPS against default SDPA, median (cuDNN-first) 0.027 (0.015) 0.014 (0.006)
  • At 1024², 10 of the 140 attention calls per UNet pass qualify: the 4096-token self-attention of the outer blocks. The 1024-token self-attention of the inner blocks and all 77-token text cross-attention are below the break-even and stay on SDPA. At 2048² the inner blocks qualify too.
  • Blind side-by-side review: 7 of 8 pairs equal, 1 judged better with SageAttention.
  • Accuracy on real activations, every call of step 0 (10 calls): cosine 1.0000, relative L2 0.003.
  • 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 SDXL. 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: the scope first opened for every UNet base, including the unmeasured SD2; it now opens for SDXL only.

Compatibility / Rollout

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

New `attention_backend: auto | sage` setting; with `sage`, eligible PyTorch SDPA calls inside a denoising scope run on SageAttention 2.2 and every other call keeps SDPA unchanged.
Each GPU's first SageAttention call is checked against SDPA and a failing GPU falls back for the session; no model family enters the scope yet.
Adds the SageAttention docs page and the regenerated config schema.
With `attention_backend: sage`, FLUX.1 denoising, including ControlNets inside the loop, enters the SageAttention scope.
RTX 4090, FLUX.1 dev FP8: 16.1 s → 14.0 s at 1024², 84.2 s → 54.8 s at 2048², peak VRAM unchanged; a blind review judged no image worse.
With `attention_backend: sage`, FLUX.2 denoising enters the SageAttention scope.
RTX 4090, FLUX.2 Klein 9B FP8: 15.6 s → 15.1 s at 1024², 27.1 s → 23.1 s at 2048², peak VRAM unchanged; a blind review judged no image worse.
With `attention_backend: sage`, Z-Image denoising enters the SageAttention scope; Z-Image Control masks its attention and keeps SDPA.
RTX 4090, Z-Image FP8: 10.4 s → 10.0 s at 1024², 33.7 s → 25.9 s at 2048², peak VRAM unchanged; a blind review judged no image worse.
With `attention_backend: sage`, Qwen-Image denoising enters the SageAttention scope.
RTX 4090, Qwen-Image 2512 GGUF Q4_K_M: 78.8 s → 73.5 s at 1024², 208.9 s → 150.3 s at 2048², peak VRAM unchanged; a blind review judged no image worse.
With `attention_backend: sage`, Krea-2 denoising enters the SageAttention scope unless INVOKE_KREA2_SDPA_BACKEND pins a PyTorch kernel.
RTX 4090, Krea-2 Turbo GGUF Q4_K_M: 11.5 s → 11.0 s at 1024², 42.7 s → 36.2 s at 2048², peak VRAM unchanged; a blind review judged no image worse.
With `attention_backend: sage`, SDXL denoising and tiled multi-diffusion enter the SageAttention scope; SD1.5 and SD2 do not.
RTX 4090: 4.4 s → 4.2 s at 1024², 19.6 s → 16.1 s at 2048², peak VRAM unchanged; a blind review judged no image worse.
@github-actions github-actions Bot added python PRs that change python files invocations PRs that change invocations backend PRs that change backend files services PRs that change app services python-tests PRs that change python tests docs PRs that change docs labels Oct 3, 2026
…ent SageAttention failures

sm86 now runs the Triton kernel upstream's sageattn picks there, and sm87 gets none; the Windows build's CUDA routing for both is unexplained and unmeasured here.
Out-of-memory and Triton cache errors fall back to SDPA for that call (three retire the device), and the first-call check runs per precision and head size.
The kernel runs with the tensors' device current, so a single-device install on cuda:1 uses SageAttention too.
An out-of-memory error now moves only the rest of that generation to SDPA, logged once, instead of counting toward retiring the device; out-of-memory is recognised by the shared is_oom_error.
Only Triton cache failures in a row retire a device, and a success resets the count.
# Conflicts:
#	invokeai/frontend/api/openapi.json

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

7.0.0 backend PRs that change backend files docs PRs that change docs invocations PRs that change invocations python PRs that change python files python-tests PRs that change python tests services PRs that change app services

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

3 participants