diff --git a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx index 4a632c09856..d4355abf977 100644 --- a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx +++ b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx @@ -31,6 +31,11 @@ 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. | + +| 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. 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. @@ -120,6 +125,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`, or in a custom node with 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() diff --git a/invokeai/app/invocations/sd/denoise_latents.py b/invokeai/app/invocations/sd/denoise_latents.py index 2a362ac6afe..0426db3978c 100644 --- a/invokeai/app/invocations/sd/denoise_latents.py +++ b/invokeai/app/invocations/sd/denoise_latents.py @@ -82,6 +82,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 @@ -1044,7 +1045,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") @@ -1198,22 +1203,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.