From 1968eb9c5862dd9e0c68509e0d6b1a878d85ce0b Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 3 Oct 2026 21:55:14 +0200 Subject: [PATCH 1/3] feat(sdxl): run SDXL attention on SageAttention when enabled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../Optimization/sage-attention.mdx | 2 + .../app/invocations/sd/denoise_latents.py | 43 +++++++++++-------- .../tiled_multi_diffusion_denoise_latents.py | 26 ++++++----- 3 files changed, 43 insertions(+), 28 deletions(-) diff --git a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx index dbbe1133297..53cace0feb9 100644 --- a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx +++ b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx @@ -31,6 +31,7 @@ compared with PyTorch's default attention there: | Z-Image (FP8) | 10.4 s → 10.0 s (−5 %) | 33.7 s → 25.9 s (−23 %) | Z-Image Control masks its attention and keeps PyTorch's. | | Qwen-Image 2512 (GGUF Q4_K_M) | 78.8 s → 73.5 s (−7 %) | 208.9 s → 150.3 s (−28 %) | One block reaches about 17 % error per call; images were still judged equivalent. | | Krea-2 Turbo (GGUF Q4_K_M) | 11.5 s → 11.0 s (−4 %) | 42.7 s → 36.2 s (−15 %) | Off while `INVOKE_KREA2_SDPA_BACKEND` pins a PyTorch kernel. | +| SDXL | 4.4 s → 4.2 s (−5 %) | 19.6 s → 16.1 s (−18 %) | Text cross-attention stays on PyTorch (too short to gain). Also serves ControlNets. | At 1024×1024 most of a step is spent outside attention, so the gain is small; at high resolutions and for video, attention dominates. GGUF models spend much of each step dequantizing weights, which SageAttention does not touch. @@ -119,6 +120,7 @@ Everything else keeps PyTorch's attention: - **Short sequences** — fewer than 1536 query or 512 key tokens. There the cost of quantizing the inputs outweighs the faster kernel. - **Head sizes other than 64 and 128.** +- **SDXL with `attention_type: sliced`** — sliced attention does not go through PyTorch's attention kernel. - **Models not listed under Supported models**, such as SD 1.5 and SD 2. Pinning a PyTorch kernel through diffusers' `DIFFUSERS_ATTN_BACKEND` does not stop SageAttention; set diff --git a/invokeai/app/invocations/sd/denoise_latents.py b/invokeai/app/invocations/sd/denoise_latents.py index aa5373fd0f8..c80e24e8a62 100644 --- a/invokeai/app/invocations/sd/denoise_latents.py +++ b/invokeai/app/invocations/sd/denoise_latents.py @@ -81,6 +81,7 @@ from invokeai.backend.util.fp8 import get_model_compute_dtype from invokeai.backend.util.hotfixes import ControlNetModel from invokeai.backend.util.mask import to_standard_float_mask +from invokeai.backend.util.sage_attention import sage_attention_scope from invokeai.backend.util.silence_warnings import SilenceWarnings @@ -1008,7 +1009,11 @@ def step_callback(state: PipelineIntermediateState) -> None: sd_backend = StableDiffusionBackend(unet, scheduler) denoise_ctx.unet = unet - result_latents = sd_backend.latents_from_embeddings(denoise_ctx, ext_manager) + # SageAttention is validated on SDXL only; SD1.5's head sizes never qualify, SD2 is unmeasured. + with ( + sage_attention_scope() if self.unet.unet.base == BaseModelType.StableDiffusionXL else nullcontext() + ): + result_latents = sd_backend.latents_from_embeddings(denoise_ctx, ext_manager) # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699 result_latents = result_latents.detach().to("cpu") @@ -1162,22 +1167,26 @@ def _lora_loader() -> Iterator[PatchSpec]: if self.hidiffusion else nullcontext() ): - result_latents = pipeline.latents_from_embeddings( - latents=latents, - timesteps=timesteps, - init_timestep=init_timestep, - noise=noise, - seed=seed, - mask=mask, - masked_latents=masked_latents, - is_gradient_mask=gradient_mask, - scheduler_step_kwargs=scheduler_step_kwargs, - conditioning_data=conditioning_data, - control_data=controlnet_data, - ip_adapter_data=ip_adapter_data, - t2i_adapter_data=t2i_adapter_data, - callback=step_callback, - ) + # SageAttention is validated on SDXL only; SD1.5's head sizes never qualify, SD2 is unmeasured. + with ( + sage_attention_scope() if self.unet.unet.base == BaseModelType.StableDiffusionXL else nullcontext() + ): + result_latents = pipeline.latents_from_embeddings( + latents=latents, + timesteps=timesteps, + init_timestep=init_timestep, + noise=noise, + seed=seed, + mask=mask, + masked_latents=masked_latents, + is_gradient_mask=gradient_mask, + scheduler_step_kwargs=scheduler_step_kwargs, + conditioning_data=conditioning_data, + control_data=controlnet_data, + ip_adapter_data=ip_adapter_data, + t2i_adapter_data=t2i_adapter_data, + callback=step_callback, + ) # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699 result_latents = result_latents.to("cpu") diff --git a/invokeai/app/invocations/sd/tiled_multi_diffusion_denoise_latents.py b/invokeai/app/invocations/sd/tiled_multi_diffusion_denoise_latents.py index 7e4e93428ab..039f5a711f4 100644 --- a/invokeai/app/invocations/sd/tiled_multi_diffusion_denoise_latents.py +++ b/invokeai/app/invocations/sd/tiled_multi_diffusion_denoise_latents.py @@ -1,5 +1,5 @@ import copy -from contextlib import ExitStack +from contextlib import ExitStack, nullcontext from typing import Iterator import torch @@ -22,6 +22,7 @@ from invokeai.app.invocations.sd.controlnet import ControlField from invokeai.app.invocations.sd.denoise_latents import DenoiseLatentsInvocation, get_scheduler from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.taxonomy import BaseModelType from invokeai.backend.patches.layer_patcher import LayerPatcher, PatchSpec from invokeai.backend.patches.model_patch_raw import ModelPatchRaw from invokeai.backend.stable_diffusion.diffusers_pipeline import ControlNetData, PipelineIntermediateState @@ -36,6 +37,7 @@ from invokeai.backend.tiles.utils import TBLR from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.fp8 import get_model_compute_dtype +from invokeai.backend.util.sage_attention import sage_attention_scope def crop_controlnet_data(control_data: ControlNetData, latent_region: TBLR) -> ControlNetData: @@ -279,16 +281,18 @@ def _lora_loader() -> Iterator[PatchSpec]: ) # Run Multi-Diffusion denoising. - result_latents = pipeline.multi_diffusion_denoise( - multi_diffusion_conditioning=multi_diffusion_conditioning, - target_overlap=latent_tile_overlap, - latents=latents, - scheduler_step_kwargs=scheduler_step_kwargs, - noise=noise, - timesteps=timesteps, - init_timestep=init_timestep, - callback=step_callback, - ) + # SageAttention is validated on SDXL only; SD1.5's head sizes never qualify, SD2 is unmeasured. + with sage_attention_scope() if self.unet.unet.base == BaseModelType.StableDiffusionXL else nullcontext(): + result_latents = pipeline.multi_diffusion_denoise( + multi_diffusion_conditioning=multi_diffusion_conditioning, + target_overlap=latent_tile_overlap, + latents=latents, + scheduler_step_kwargs=scheduler_step_kwargs, + noise=noise, + timesteps=timesteps, + init_timestep=init_timestep, + callback=step_callback, + ) result_latents = result_latents.to("cpu") # TODO(ryand): I copied this from DenoiseLatentsInvocation. I'm not sure if it's actually important. From 9507eccfb65b5bc4cf5b1ff70b654fc7fe140ec9 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 3 Oct 2026 21:55:15 +0200 Subject: [PATCH 2/3] feat(ltx2): run LTX-2 video self-attention on SageAttention when enabled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With `attention_backend: sage`, LTX-2 denoising enters the SageAttention scope; video self-attention qualifies, while the masked text cross-attention and the short audio attention keep SDPA. RTX 4090, LTX-2.5 distilled int8 at 1248×704, 121 frames: 72.0 s → 63.4 s, peak VRAM unchanged, no added flicker. --- .../Optimization/sage-attention.mdx | 4 +++ invokeai/app/invocations/ltx2/ltx2_denoise.py | 26 ++++++++++--------- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx index 53cace0feb9..7b21d964d13 100644 --- a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx +++ b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx @@ -33,6 +33,10 @@ compared with PyTorch's default attention there: | Krea-2 Turbo (GGUF Q4_K_M) | 11.5 s → 11.0 s (−4 %) | 42.7 s → 36.2 s (−15 %) | Off while `INVOKE_KREA2_SDPA_BACKEND` pins a PyTorch kernel. | | SDXL | 4.4 s → 4.2 s (−5 %) | 19.6 s → 16.1 s (−18 %) | Text cross-attention stays on PyTorch (too short to gain). Also serves ControlNets. | +| Video model (as measured) | Resolution | Generation | Notes | +| --- | --- | --- | --- | +| LTX-2.5 distilled (int8) | 1248×704, 121 frames | 72.0 s → 63.4 s (−12 %) | Video self-attention only; text cross-attention and the audio streams stay on PyTorch. | + At 1024×1024 most of a step is spent outside attention, so the gain is small; at high resolutions and for video, attention dominates. GGUF models spend much of each step dequantizing weights, which SageAttention does not touch. Windows' PyTorch has no flash-attention kernel; on Linux PyTorch's own baseline is faster and the gain correspondingly diff --git a/invokeai/app/invocations/ltx2/ltx2_denoise.py b/invokeai/app/invocations/ltx2/ltx2_denoise.py index 87876e50959..740aeac340c 100644 --- a/invokeai/app/invocations/ltx2/ltx2_denoise.py +++ b/invokeai/app/invocations/ltx2/ltx2_denoise.py @@ -59,6 +59,7 @@ from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import LTX2ConditioningInfo from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.sage_attention import sage_attention_scope LTX2Schedule = Literal["auto", "dev", "distilled"] @@ -616,18 +617,19 @@ def step_callback(step: int, total_steps: int, video_x0: torch.Tensor) -> None: ) try: - video_latents, audio_latents = denoise( - transformer=transformer, - state=state, - positive=positive, - negative=negative, - guidance=guidance, - fps=self.fps, - dtype=inference_dtype, - device=device, - step_callback=step_callback, - is_canceled=context.util.is_canceled, - ) + with sage_attention_scope(): + video_latents, audio_latents = denoise( + transformer=transformer, + state=state, + positive=positive, + negative=negative, + guidance=guidance, + fps=self.fps, + dtype=inference_dtype, + device=device, + step_callback=step_callback, + is_canceled=context.util.is_canceled, + ) finally: progress.close() From b3c3e4300a053abf02d5956e7c67fa12af3cc5b4 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 12 Oct 2026 05:14:06 +0200 Subject: [PATCH 3/3] docs(attention): state what SageAttention does to LTX-2's audio The audio reads the video through cross-attention; against PyTorch's attention it changed no more than an exact kernel swap does, with no shift against the picture and the same loudness (8 clips, fixed seeds, clap and speech prompts included). --- .../content/docs/configuration/Optimization/sage-attention.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx index 3739e01f9f7..8d45dcf38f5 100644 --- a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx +++ b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx @@ -35,7 +35,7 @@ compared with PyTorch's default attention there: | Video model (as measured) | Resolution | Generation | Notes | | --- | --- | --- | --- | -| LTX-2.5 distilled (int8) | 1248×704, 121 frames | 72.0 s → 63.4 s (−12 %) | Video self-attention only; text cross-attention and the audio streams stay on PyTorch. | +| LTX-2.5 distilled (int8) | 1248×704, 121 frames | 72.0 s → 63.4 s (−12 %) | Video self-attention only; text cross-attention and the audio streams stay on PyTorch. The audio reads the video through cross-attention, so it changes too, but measured no more than with an exact kernel swap: no shift against the picture, same loudness. | At 1024×1024 most of a step is spent outside attention, so the gain is small; at high resolutions and for video, attention dominates. GGUF models spend much of each step dequantizing weights, which SageAttention does not touch.