feat(ltx2): run LTX-2 video self-attention on SageAttention when enabled

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

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()
 

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")

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.