diff --git a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx
new file mode 100644
index 00000000000..f6e6dbe2567
--- /dev/null
+++ b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx
@@ -0,0 +1,148 @@
+---
+title: SageAttention
+sidebar:
+ order: 5
+---
+
+import { Steps } from '@astrojs/starlight/components';
+
+[SageAttention](https://github.com/thu-ml/SageAttention) is a faster attention kernel for NVIDIA GPUs. It quantizes the
+attention inputs to 8 bits inside the kernel. On an RTX 4090 under Windows that makes the attention of the image and
+video models two to three times faster than PyTorch's own kernels. How much of a generation that saves depends on how
+much of each step is attention: it grows with resolution and video length.
+
+It is **opt-in**. Because the kernel is quantized, an image made with SageAttention is not pixel-identical to one made
+without it at the same seed — composition details can shift the way they do between two GPU models. In a blind
+side-by-side comparison of 48 image pairs from six image model families, no pair was judged worse with SageAttention;
+the video models were checked against PyTorch's attention for picture changes and frame-to-frame flicker.
+
+## Supported models
+
+SageAttention is enabled per model family once that family has been measured with it. All other models keep PyTorch's
+attention.
+
+Whole generation (text encoding, denoising and decoding) with the model already loaded, on an RTX 4090 under Windows,
+compared with PyTorch's default attention there:
+
+| Model (as measured) | 1024×1024 | 2048×2048 | Notes |
+| --- | --- | --- | --- |
+| FLUX.1 dev (FP8) | 16.1 s → 14.0 s (−13 %) | 84.2 s → 54.8 s (−35 %) | Also serves ControlNets inside the denoising loop. |
+| FLUX.2 Klein 9B (FP8) | 15.6 s → 15.1 s (−3 %) | 27.1 s → 23.1 s (−15 %) | |
+| 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.
+Windows' PyTorch has no flash-attention kernel; on Linux PyTorch's own baseline is faster and the gain correspondingly
+smaller (not measured).
+
+## Requirements
+
+- **An NVIDIA GPU with compute capability 8.0 or newer** — RTX 30-series, RTX 40-series, RTX 50-series, A100, H100.
+ InvokeAI has measured it on an RTX 4090; the other generations use the kernels SageAttention picks for them, which
+ InvokeAI has not validated (on RTX 30-series cards, its Triton kernel). Jetson Orin is not supported.
+- **SageAttention 2.2, installed separately.** InvokeAI was validated with 2.2.0. The `sageattention` package on PyPI is
+ the older SageAttention 1 and is not used. On Windows it also needs the `triton-windows` package.
+- Not available on AMD (ROCm), Apple Silicon (MPS), Intel (XPU) or CPU. There InvokeAI keeps PyTorch's attention.
+
+## Installing
+
+SageAttention goes into the Python environment InvokeAI runs from — the `.venv` folder in your install directory. The
+commands use [uv](https://docs.astral.sh/uv/), which InvokeAI's environment is managed with; install it first if
+`uv --version` does not work in your terminal.
+
+### Windows
+
+
+1. Open a terminal in your InvokeAI install directory.
+2. Install Triton for Windows, matching the Triton version of InvokeAI's PyTorch (Triton 3.7 for PyTorch 2.13):
+
+ ```powershell
+ uv pip install --python .venv\Scripts\python.exe "triton-windows>=3.7,<3.8"
+ ```
+
+3. Install a SageAttention 2.2 wheel from the [SageAttention for Windows releases](https://github.com/woct0rdho/SageAttention/releases)
+ that matches InvokeAI's PyTorch (CUDA 13, PyTorch 2.10 or newer), without letting it touch PyTorch:
+
+ ```powershell
+ uv pip install --python .venv\Scripts\python.exe --no-deps "https://github.com/woct0rdho/SageAttention/releases/download/v2.2.0-windows.post6/sageattention-2.2.0%2Bcu130torch2.10.0andhigher.post6-cp310-abi3-win_amd64.whl"
+ ```
+
+
+### Linux
+
+SageAttention 2 has no prebuilt wheels for Linux. Build it from source into InvokeAI's environment
+(`.venv/bin/python`) as described in [SageAttention's installation instructions](https://github.com/thu-ml/SageAttention#installation),
+with the CUDA toolkit that matches InvokeAI's PyTorch (CUDA 13) installed. InvokeAI's environment has no `pip`; use
+`uv pip install --python .venv/bin/python ...` where those instructions say `pip install`. Triton ships with PyTorch on
+Linux.
+
+:::caution[Keep it matched to InvokeAI's PyTorch]
+SageAttention is built against PyTorch, and Triton must match PyTorch's version. When an update moves InvokeAI to a new
+PyTorch, install the matching `triton-windows` again on Windows, and on Linux rebuild SageAttention. A wheel for a
+different CUDA major version, or a repair that re-creates the environment, needs a matching SageAttention too. Until
+then InvokeAI logs why it cannot use it and generates with PyTorch's attention.
+:::
+
+## Enabling
+
+Set the backend in your [`invokeai.yaml`](/configuration/invokeai-yaml/) and restart InvokeAI:
+
+```yaml
+attention_backend: sage
+```
+
+The startup log confirms it:
+
+```
+SageAttention 2.2.0+cu130torch2.10.0andhigher.post6 enabled for the attention inside supported diffusion models ...
+```
+
+and the first generation that uses it names the GPU and the kernel:
+
+```
+SageAttention 2.2.0+cu130torch2.10.0andhigher.post6 serves diffusion-model attention on cuda:0 (NVIDIA GeForce RTX 4090, sm_89, sageattn_qk_int8_pv_fp8_cuda).
+```
+
+If SageAttention cannot be used, the startup log says why — package missing, Triton missing, GPU too old — and
+InvokeAI generates with PyTorch's attention as usual. With `log_level: debug`, every generation logs how many attention
+calls ran on SageAttention and why the others did not. To go back, set `attention_backend: auto`.
+
+## What stays on PyTorch's attention
+
+SageAttention serves the attention inside the denoising model of the families under [Supported models](#supported-models).
+Everything else keeps PyTorch's attention:
+
+- **Text encoders, image encoders and VAEs.** They run outside the denoising loop. What runs inside it, such as a
+ ControlNet, is served like the model itself.
+- **Masked attention**, such as regional prompts. SageAttention has no mask input.
+- **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
+`torch.nn.attention.sdpa_kernel`, does not stop SageAttention; set `attention_backend: auto` to measure PyTorch's own
+kernels.
+
+## Notes
+
+- **The first generation pauses briefly.** SageAttention compiles a few Triton kernels on first use, about two seconds
+ on an RTX 4090; they are cached on disk afterwards.
+- **Each GPU's first SageAttention call is checked against PyTorch's attention**, once for each precision and head
+ size. If the result is wrong or the kernel fails, InvokeAI logs a warning once and that GPU uses PyTorch's attention
+ until restart. Some GPU faults (an illegal memory access, for example) leave the GPU unusable for the rest of the
+ process: then PyTorch's attention fails the same way, and only a restart helps. A fault the GPU raises later, while
+ running a kernel, cannot be caught at all and ends the generation.
+- **Memory.** Peak VRAM was unchanged for every measured model. Each call holds quantized copies of its inputs, about
+ 1.7 bytes per element of the query, key and value together — around 70 MB for FLUX.1 at 1024×1024, and around 2 GB
+ per call for Wan 2.2 14B at 1280×720 with 81 frames (not measured). When a call does not fit, the rest of that
+ generation runs on PyTorch's attention, which needs less, and the log says so; the next generation tries
+ SageAttention again.
+- **Recall.** `attention_backend` is not stored in an image's metadata, so an image recalled after switching it is
+ generated slightly differently.
+- **Black or noisy images** have been reported with SageAttention for some models in other applications. None appeared
+ in InvokeAI's tests; if you see one, set `attention_backend: auto` and report it with the model and settings.
diff --git a/docs/src/content/docs/troubleshooting/faq.mdx b/docs/src/content/docs/troubleshooting/faq.mdx
index bdd339130bb..5b95ff6882a 100644
--- a/docs/src/content/docs/troubleshooting/faq.mdx
+++ b/docs/src/content/docs/troubleshooting/faq.mdx
@@ -44,7 +44,7 @@ Follow the same steps to scan and import the missing models.
## Triton error on startup
-This can be safely ignored. Invoke doesn't use Triton, but if you are on Linux and wish to dismiss the error, you can install Triton.
+This can be safely ignored. Invoke only uses Triton for [SageAttention](/configuration/optimization/sage-attention/), which is off unless you enable it. If you are on Linux and wish to dismiss the error, you can install Triton.
## Unable to Copy on Firefox
diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json
index 8a1cc4fb1ba..15e4a532fb6 100644
--- a/docs/src/generated/settings.json
+++ b/docs/src/generated/settings.json
@@ -798,6 +798,20 @@
"type": "typing.Literal['auto', 'balanced', 'max', 1, 2, 3, 4, 5, 6, 7, 8]",
"validation": {}
},
+ {
+ "category": "GENERATION",
+ "default": "auto",
+ "description": "Attention kernel for the attention inside diffusion models. `auto` lets PyTorch choose among its exact SDPA kernels and never selects a quantized one. `sage` uses SageAttention 2 for the models listed in the SageAttention docs: faster, most at high resolutions and for video, but quantized, so images differ slightly from `auto` at the same seed. SageAttention is installed separately and needs an NVIDIA GPU with compute capability 8.0 or newer; masked attention, text encoders and VAEs keep PyTorch SDPA.",
+ "env_var": "INVOKEAI_ATTENTION_BACKEND",
+ "literal_values": [
+ "auto",
+ "sage"
+ ],
+ "name": "attention_backend",
+ "required": false,
+ "type": "typing.Literal['auto', 'sage']",
+ "validation": {}
+ },
{
"category": "GENERATION",
"default": false,
diff --git a/invokeai/app/invocations/flux/flux_denoise.py b/invokeai/app/invocations/flux/flux_denoise.py
index 19f0924c940..c1aa6a3d691 100644
--- a/invokeai/app/invocations/flux/flux_denoise.py
+++ b/invokeai/app/invocations/flux/flux_denoise.py
@@ -69,6 +69,7 @@
from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo
from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.sage_attention import sage_attention_scope
@invocation(
@@ -525,26 +526,27 @@ def _run_diffusion(
else:
context.logger.debug(f"DyPE disabled: resolution={self.width}x{self.height}, preset={self.dype_preset}")
- x = denoise(
- model=transformer,
- img=x,
- img_ids=img_ids,
- pos_regional_prompting_extension=pos_regional_prompting_extension,
- neg_regional_prompting_extension=neg_regional_prompting_extension,
- timesteps=timesteps,
- step_callback=self._build_step_callback(context),
- guidance=self.guidance,
- cfg_scale=cfg_scale,
- inpaint_extension=inpaint_extension,
- controlnet_extensions=controlnet_extensions,
- pos_ip_adapter_extensions=pos_ip_adapter_extensions,
- neg_ip_adapter_extensions=neg_ip_adapter_extensions,
- img_cond=img_cond,
- img_cond_seq=img_cond_seq,
- img_cond_seq_ids=img_cond_seq_ids,
- dype_extension=dype_extension,
- scheduler=scheduler,
- )
+ with sage_attention_scope():
+ x = denoise(
+ model=transformer,
+ img=x,
+ img_ids=img_ids,
+ pos_regional_prompting_extension=pos_regional_prompting_extension,
+ neg_regional_prompting_extension=neg_regional_prompting_extension,
+ timesteps=timesteps,
+ step_callback=self._build_step_callback(context),
+ guidance=self.guidance,
+ cfg_scale=cfg_scale,
+ inpaint_extension=inpaint_extension,
+ controlnet_extensions=controlnet_extensions,
+ pos_ip_adapter_extensions=pos_ip_adapter_extensions,
+ neg_ip_adapter_extensions=neg_ip_adapter_extensions,
+ img_cond=img_cond,
+ img_cond_seq=img_cond_seq,
+ img_cond_seq_ids=img_cond_seq_ids,
+ dype_extension=dype_extension,
+ scheduler=scheduler,
+ )
x = unpack(x.float(), self.height, self.width)
return x
diff --git a/invokeai/app/invocations/flux2/flux2_denoise.py b/invokeai/app/invocations/flux2/flux2_denoise.py
index a874c8101f0..19e4fd48738 100644
--- a/invokeai/app/invocations/flux2/flux2_denoise.py
+++ b/invokeai/app/invocations/flux2/flux2_denoise.py
@@ -59,6 +59,7 @@
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo
from invokeai.backend.util.attention import sdpa_score_matrix_bytes
from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.sage_attention import sage_attention_scope
# FLUX.2 attention geometry. The head dim is 128 across every variant and the head count follows
# the hidden size (Klein 4B: 3072/24, Klein 9B: 4096/32, [dev] 6144/48), so the width is the single
@@ -623,25 +624,26 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
"Regional masks will be ignored for this generation."
)
- x = denoise(
- model=transformer,
- img=x,
- img_ids=img_ids,
- txt=txt,
- txt_ids=txt_ids,
- timesteps=timesteps,
- step_callback=self._build_step_callback(context),
- guidance=self.guidance,
- cfg_scale=cfg_scale_list,
- neg_txt=neg_txt,
- neg_txt_ids=neg_txt_ids,
- scheduler=scheduler,
- mu=mu,
- inpaint_extension=inpaint_extension,
- img_cond_seq=img_cond_seq,
- img_cond_seq_ids=img_cond_seq_ids,
- pos_joint_attention_kwargs=pos_joint_attention_kwargs,
- )
+ with sage_attention_scope():
+ x = denoise(
+ model=transformer,
+ img=x,
+ img_ids=img_ids,
+ txt=txt,
+ txt_ids=txt_ids,
+ timesteps=timesteps,
+ step_callback=self._build_step_callback(context),
+ guidance=self.guidance,
+ cfg_scale=cfg_scale_list,
+ neg_txt=neg_txt,
+ neg_txt_ids=neg_txt_ids,
+ scheduler=scheduler,
+ mu=mu,
+ inpaint_extension=inpaint_extension,
+ img_cond_seq=img_cond_seq,
+ img_cond_seq_ids=img_cond_seq_ids,
+ pos_joint_attention_kwargs=pos_joint_attention_kwargs,
+ )
# Apply BN denormalization if BN stats are available
# The diffusers Flux2KleinPipeline applies: latents = latents * bn_std + bn_mean
diff --git a/invokeai/app/invocations/krea2/krea2_denoise.py b/invokeai/app/invocations/krea2/krea2_denoise.py
index 53874e4e11a..294fd018397 100644
--- a/invokeai/app/invocations/krea2/krea2_denoise.py
+++ b/invokeai/app/invocations/krea2/krea2_denoise.py
@@ -63,6 +63,7 @@
from invokeai.backend.util.attention import sdpa_score_matrix_bytes
from invokeai.backend.util.devices import TorchDevice
from invokeai.backend.util.logging import InvokeAILogger
+from invokeai.backend.util.sage_attention import sage_attention_scope
# Krea-2 latent channels (Qwen-Image VAE z_dim). The packed transformer in_channels is 16 * patch_size**2 = 64.
KREA2_LATENT_CHANNELS = 16
@@ -625,6 +626,11 @@ def _run_diffusion(self, context: InvocationContext):
# the sampler will use.
style_extension.prepare([sigma.item() for sigma in sigmas_sched[:total_steps]])
+ # An INVOKE_KREA2_SDPA_BACKEND override pins one PyTorch kernel to measure it; SageAttention must not
+ # serve those calls instead.
+ if resolve_krea2_sdpa_backends().override is None:
+ exit_stack.enter_context(sage_attention_scope())
+
benchmark = _Krea2StepBenchmark.create(device)
for step_idx, t in enumerate(tqdm(timesteps_sched)):
diff --git a/invokeai/app/invocations/qwen_image/qwen_image_denoise.py b/invokeai/app/invocations/qwen_image/qwen_image_denoise.py
index a0b46234220..2f63e215783 100644
--- a/invokeai/app/invocations/qwen_image/qwen_image_denoise.py
+++ b/invokeai/app/invocations/qwen_image/qwen_image_denoise.py
@@ -34,6 +34,7 @@
from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import QwenImageConditioningInfo
from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.sage_attention import sage_attention_scope
@invocation(
@@ -485,6 +486,8 @@ def _run_diffusion(self, context: InvocationContext):
)
)
+ exit_stack.enter_context(sage_attention_scope())
+
for step_idx, t in enumerate(tqdm(timesteps_sched)):
# The pipeline passes timestep / 1000 to the transformer
timestep = t.expand(latents.shape[0]).to(inference_dtype)
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.
diff --git a/invokeai/app/invocations/z_image/z_image_denoise.py b/invokeai/app/invocations/z_image/z_image_denoise.py
index 66623c30586..f71847fcb91 100644
--- a/invokeai/app/invocations/z_image/z_image_denoise.py
+++ b/invokeai/app/invocations/z_image/z_image_denoise.py
@@ -39,6 +39,7 @@
from invokeai.backend.util.attention import sdpa_score_matrix_bytes
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
from invokeai.backend.z_image.extensions.regional_prompting_extension import ZImageRegionalPromptingExtension
from invokeai.backend.z_image.text_conditioning import ZImageTextConditioning
from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter
@@ -700,6 +701,8 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
)
)
+ exit_stack.enter_context(sage_attention_scope())
+
# Denoising loop - supports both built-in Euler and diffusers schedulers
# Track user-facing step for progress (accounts for Heun's double steps)
user_step = 0
diff --git a/invokeai/app/run_app.py b/invokeai/app/run_app.py
index d838dba5eb5..a0bf2d97801 100644
--- a/invokeai/app/run_app.py
+++ b/invokeai/app/run_app.py
@@ -116,7 +116,7 @@ def run_app() -> None:
logger.warning(f"Port {orig_config_port} is already in use. Using port {app_config.port}.")
# Miscellaneous startup tasks.
- apply_monkeypatches()
+ apply_monkeypatches(app_config.attention_backend)
register_mime_types()
check_cudnn(logger)
check_cuda_build_compatibility(logger)
diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py
index 90cf75be8f3..7ef59bd3f0d 100644
--- a/invokeai/app/services/config/config_default.py
+++ b/invokeai/app/services/config/config_default.py
@@ -28,6 +28,7 @@
LEGACY_INIT_FILE = Path("invokeai.init")
PRECISION = Literal["auto", "float16", "bfloat16", "float32"]
ATTENTION_TYPE = Literal["auto", "normal", "xformers", "sliced", "torch-sdp"]
+ATTENTION_BACKEND = Literal["auto", "sage"]
NOISE_DTYPE = Literal["float32", "float16"]
ROCM_AOTRITON = Literal["auto", "on", "off"]
ATTENTION_SLICE_SIZE = Literal["auto", "balanced", "max", 1, 2, 3, 4, 5, 6, 7, 8]
@@ -135,6 +136,7 @@ class InvokeAIAppConfig(BaseSettings):
pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.
attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`
attention_slice_size: Slice size, valid when attention_type=="sliced".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`
+ attention_backend: Attention kernel for the attention inside diffusion models. `auto` lets PyTorch choose among its exact SDPA kernels and never selects a quantized one. `sage` uses SageAttention 2 for the models listed in the SageAttention docs: faster, most at high resolutions and for video, but quantized, so images differ slightly from `auto` at the same seed. SageAttention is installed separately and needs an NVIDIA GPU with compute capability 8.0 or newer; masked attention, text encoders and VAEs keep PyTorch SDPA.
Valid values: `auto`, `sage`
force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.
auto_tiled_decode: Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.
pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.
@@ -269,6 +271,7 @@ class InvokeAIAppConfig(BaseSettings):
pid_memory_optimization: bool = Field(default=False, description="Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.")
attention_type: ATTENTION_TYPE = Field(default="auto", description="Attention type.")
attention_slice_size: ATTENTION_SLICE_SIZE = Field(default="auto", description='Slice size, valid when attention_type=="sliced".')
+ attention_backend: ATTENTION_BACKEND = Field(default="auto", description="Attention kernel for the attention inside diffusion models. `auto` lets PyTorch choose among its exact SDPA kernels and never selects a quantized one. `sage` uses SageAttention 2 for the models listed in the SageAttention docs: faster, most at high resolutions and for video, but quantized, so images differ slightly from `auto` at the same seed. SageAttention is installed separately and needs an NVIDIA GPU with compute capability 8.0 or newer; masked attention, text encoders and VAEs keep PyTorch SDPA.")
force_tiled_decode: bool = Field(default=False, description="Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.")
auto_tiled_decode: bool = Field(default=True, description="Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.")
pil_compress_level: int = Field(default=1, description="The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.")
diff --git a/invokeai/app/util/startup_utils.py b/invokeai/app/util/startup_utils.py
index aa88f104b5e..886872790c2 100644
--- a/invokeai/app/util/startup_utils.py
+++ b/invokeai/app/util/startup_utils.py
@@ -8,6 +8,8 @@
import torch
+from invokeai.app.services.config.config_default import ATTENTION_BACKEND
+
def find_open_port(port: int) -> int:
"""Find a port not in use starting at given port"""
@@ -200,8 +202,8 @@ def enable_dev_reload(custom_nodes_path=None) -> None:
jurigged.watch(pattern=paths, logger=InvokeAILogger.get_logger(name="jurigged").info)
-def apply_monkeypatches() -> None:
- """Apply monkeypatches to fix issues with third-party libraries."""
+def apply_monkeypatches(attention_backend: ATTENTION_BACKEND = "auto") -> None:
+ """Apply monkeypatches to fix issues with third-party libraries, and install the configured attention backend."""
import invokeai.backend.util.hotfixes # noqa: F401 (monkeypatching on import)
@@ -220,6 +222,12 @@ def apply_monkeypatches() -> None:
install_rocm_sdpa_guard()
+ # Opt-in SageAttention for diffusion-model attention. Installs nothing unless asked for, and never on ROCm.
+ if attention_backend == "sage":
+ from invokeai.backend.util.sage_attention import install_sage_attention
+
+ install_sage_attention()
+
def register_mime_types() -> None:
"""Register additional mime types for windows."""
diff --git a/invokeai/backend/util/sage_attention.py b/invokeai/backend/util/sage_attention.py
new file mode 100644
index 00000000000..7510e5a278f
--- /dev/null
+++ b/invokeai/backend/util/sage_attention.py
@@ -0,0 +1,429 @@
+"""SageAttention for the attention inside diffusion models, opt-in via ``attention_backend: sage``.
+
+SageAttention (thu-ml) quantizes Q/K to INT8 and P/V to FP8 or FP16 inside the attention kernel. On an RTX 4090 under
+Windows it runs the attention shapes of the image and video transformers 2-3x faster than PyTorch's fastest SDPA kernel
+there, and the gap widens with the token count. It is lossy: images differ from SDPA at the same seed (in a blind
+comparison of 48 pairs over six architectures, none worse). That is why it is opt-in, and why it only runs where a
+denoise invocation asks for it.
+
+How it is wired:
+
+* ``install_sage_attention`` rebinds ``torch.nn.functional.scaled_dot_product_attention`` once at startup, like
+ ``install_rocm_sdpa_guard``. Every caller that resolves the function at call time inherits it: InvokeAI's own
+ transformer code, diffusers' attention processors and diffusers' ``native`` dispatch backend. diffusers' own
+ ``sage`` backend is not used: it raises on any mask, is process-global, and never sees the first two routes.
+* The wrapper does nothing outside ``sage_attention_scope``. Text encoders, image encoders and VAEs therefore keep
+ PyTorch's kernels, as do models whose denoise invocation has not opted in. The scope is a ``ContextVar``, so it is
+ per thread: InvokeAI runs one generation session per GPU concurrently, and a process-global switch would leak
+ between them (see ``sdpa_scope``).
+* Inside the scope, a call goes to SageAttention only when ``_ineligible`` finds nothing wrong with it -- no mask,
+ a head size the kernels run natively, half precision, sequences long enough to be worth the quantization.
+ Everything else falls back to PyTorch with its arguments untouched.
+"""
+
+import functools
+import inspect
+import math
+import re
+import threading
+from collections import Counter
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+from invokeai.backend.util.logging import InvokeAILogger
+from invokeai.backend.util.oom import is_oom_error
+
+logger = InvokeAILogger.get_logger(__name__)
+
+_SENTINEL = "_invokeai_sage_attention"
+_MIN_VERSION = (2, 2, 0)
+# Head sizes the kernels run natively. Other sizes are zero-padded by SageAttention (40 -> 64, 80 -> 128), which costs
+# copies and measured slower than SDPA at the shapes that have them (SD1.5).
+_HEAD_DIMS = frozenset({64, 128})
+# Shortest query and key sequences SageAttention is used for. It quantizes Q, K and V in separate kernels first, a cost
+# that grows with each sequence's length, while what it saves grows with their product -- so a short side on either
+# end loses. Measured on an RTX 4090 against the faster of SDPA's default and cuDNN kernels: keys of 77 or 128 tokens
+# run 0.2-0.8x (SDXL cross-attention, LTX-2's audio streams), 256 break even, 512 and longer run 1.4-1.9x; queries of
+# 512 or 1024 tokens against 13.7k keys run 0.6x and 1.0x, 1536 and longer 1.4x and up.
+_MIN_QUERY_LEN = 1536
+_MIN_KEY_LEN = 512
+_HALF_PRECISION = (torch.float16, torch.bfloat16)
+# Largest relative L2 difference from PyTorch SDPA a device's first SageAttention result may have. Real attention
+# stays well below it (worst measured: 0.19, in Qwen-Image's 53rd block); a kernel that never ran leaves uninitialized
+# memory, which lands near 1 or is not finite.
+_FIRST_CALL_MAX_ERROR = 0.5
+# Query rows of the first call that are recomputed with PyTorch for that comparison, from the end so the last, partial
+# tile is among them. A kernel that never ran garbles every row, so a slice catches what the whole output would; the
+# whole output costs ~14 bytes per element to compare -- about 5 GB for one call of Wan 2.2 14B at 720p.
+_FIRST_CALL_ROWS = 1024
+# Triton failing to write a compiled kernel to its cache (on Windows, when two GPUs compile the same kernel at once)
+# says nothing about the kernel: that call falls back and the next tries again. A device that fails this way this many
+# times in a row is retired like any other, so a cache that can never be written does not recompile on every call.
+_MAX_CACHE_FAILURES = 3
+
+DOCS_URL = "https://invoke-ai.github.io/InvokeAI-7/configuration/optimization/sage-attention/"
+
+
+@dataclass
+class _Usage:
+ """What one scope's attention calls ran on. Logged at DEBUG when the scope closes."""
+
+ served: int = 0
+ fell_back: Counter[str] = field(default_factory=Counter)
+ # Set once SageAttention's quantized copies did not fit; the rest of the scope then stays on SDPA.
+ memory_short: bool = False
+
+
+_scope: ContextVar[_Usage | None] = ContextVar("invokeai_sage_attention_scope", default=None)
+_lock = threading.Lock()
+_disabled_devices: set[int] = set()
+# (device, dtype, head size): the kernel variants whose first result was compared with SDPA.
+_validated: set[tuple[int, torch.dtype, int]] = set()
+_cache_failures: Counter[int] = Counter()
+_announced_devices: set[int] = set()
+
+
+@contextmanager
+def sage_attention_scope() -> Iterator[None]:
+ """Let SageAttention serve the eligible attention calls this thread makes inside the block.
+
+ A no-op unless ``install_sage_attention`` installed the wrapper. Denoise invocations enter it around the
+ transformer forward passes of architectures whose output was validated against SDPA.
+ """
+ usage = _Usage()
+ token = _scope.set(usage)
+ try:
+ yield
+ finally:
+ _scope.reset(token)
+ if usage.served or usage.fell_back:
+ fallbacks = ", ".join(f"{reason}={count}" for reason, count in usage.fell_back.most_common())
+ logger.debug(
+ f"SageAttention: {usage.served + sum(usage.fell_back.values())} attention calls, "
+ f"{usage.served} on SageAttention" + (f"; on PyTorch SDPA: {fallbacks}" if fallbacks else "")
+ )
+
+
+def _import_sageattention() -> tuple[str, Any]:
+ """The only place that touches the optional package. Imported here, never at module import time."""
+ from importlib.metadata import version
+
+ import sageattention
+
+ return version("sageattention"), sageattention
+
+
+def _parse_version(version: str) -> tuple[int, int, int] | None:
+ # Wheels carry local labels: "2.2.0+cu130torch2.10.0andhigher.post6".
+ match = re.match(r"(\d+)\.(\d+)(?:\.(\d+))?", version)
+ return (int(match[1]), int(match[2]), int(match[3] or 0)) if match else None
+
+
+def _kernel(
+ sageattention: Any, capability: tuple[int, int], cuda: tuple[int, int]
+) -> Callable[..., torch.Tensor] | None:
+ """SageAttention's kernel for a GPU of this compute capability, or None where it has none.
+
+ The kernels and accumulation upstream SageAttention 2.2's ``sageattn`` (thu-ml) dispatches to: the FP16 CUDA
+ kernel on sm80, the Triton one on sm86, nothing on sm87. The Windows build moved sm86 and sm87 to the CUDA kernel
+ without saying why; neither has been measured here, so upstream's choice stands. From the Windows build come the
+ sm100 kernel and the choice by the CUDA version torch was built for: the FP16 accumulation inside the FP8 kernels
+ needs CUDA 12.8 to compile, and below that the build replaces it with a trap. Every kernel is called with
+ ``smooth_k`` off, which ``sageattn`` always leaves on and has no argument to change.
+
+ Smoothing subtracts K's mean over the sequence before quantizing; in the first block of Qwen-Image and Krea-2,
+ where one channel of Q and K sits near +600 for every image token, that leaves the INT8 product Q.K no precision for
+ the rest: cosine 0.18 and 0.53 against a float32 reference, 0.997 and 0.988 with smoothing off. Everywhere else
+ measured it changes nothing, and off is 2-11 % faster (RTX 4090). Only sm_89 has been measured in InvokeAI; the
+ sm86 Triton kernel was checked for accuracy on it.
+ """
+ major, minor = capability
+ recent_cuda = cuda >= (12, 8)
+ if (major, minor) == (8, 0): # A100: FP16 P.V
+ kernel, options = sageattention.sageattn_qk_int8_pv_fp16_cuda, {"pv_accum_dtype": "fp32"}
+ elif (major, minor) == (8, 6): # RTX 30-series, A10, A40: FP16 P.V in Triton
+ kernel, options = sageattention.sageattn_qk_int8_pv_fp16_triton, {}
+ elif (major, minor) == (8, 9): # Ada: FP8 P.V
+ kernel = sageattention.sageattn_qk_int8_pv_fp8_cuda
+ options = {"pv_accum_dtype": "fp32+fp16" if recent_cuda else "fp32+fp32"}
+ elif major == 9: # Hopper
+ kernel, options = sageattention.sageattn_qk_int8_pv_fp8_cuda_sm90, {"pv_accum_dtype": "fp32+fp32"}
+ elif (major, minor) in ((10, 0), (12, 0), (12, 1)): # Blackwell
+ kernel = sageattention.sageattn_qk_int8_pv_fp8_cuda
+ options = {"qk_quant_gran": "per_warp", "pv_accum_dtype": "fp32+fp16" if recent_cuda else "fp32"}
+ else:
+ return None
+ return functools.partial(kernel, smooth_k=False, **options)
+
+
+def _unaccepted_arguments(kernel: functools.partial) -> list[str]:
+ """Arguments this kernel would silently swallow: every SageAttention kernel ends in ``**kwargs``, so a renamed
+ ``smooth_k`` would be ignored rather than rejected, and smoothing would quietly come back."""
+ parameters = inspect.signature(kernel.func).parameters
+ return [name for name in (*kernel.keywords, "tensor_layout", "is_causal", "sm_scale") if name not in parameters]
+
+
+def _cuda_index(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> int | None:
+ """The CUDA index of q/k/v when all three live on the same CUDA device, else None.
+
+ SageAttention launches its kernels on the thread's current device without switching to the tensors' device, so the
+ wrapper makes that device current for the call. Session workers pin their thread to their device already; a
+ single-device install with `device: cuda:1` does not.
+ """
+ device = query.device
+ if device.type != "cuda" or key.device != device or value.device != device:
+ return None
+ return device.index
+
+
+def _ineligible(
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_mask: torch.Tensor | None,
+ dropout_p: float,
+ is_causal: bool,
+ enable_gqa: bool,
+ extra_arguments: bool,
+) -> str | None:
+ """Why this call must stay on PyTorch's SDPA, or None if SageAttention may serve it.
+
+ Reads tensor metadata only: anything that synchronizes the device (inspecting a mask's values, say) would cost
+ more than the kernel saves. Which device the call runs on is checked by the caller, which needs the index anyway.
+ """
+ if extra_arguments:
+ return "arguments"
+ if attn_mask is not None:
+ return "mask"
+ if dropout_p != 0.0:
+ return "dropout"
+ if is_causal:
+ # No denoiser in scope attends causally, and torch and SageAttention have not been checked to align a causal
+ # mask the same way when query and key lengths differ.
+ return "causal"
+ if query.dim() != 4 or key.dim() != 4 or value.dim() != 4 or query.is_nested:
+ return "shape"
+ if query.dtype not in _HALF_PRECISION or key.dtype != query.dtype or value.dtype != query.dtype:
+ return "dtype"
+ batch, heads, query_len, head_dim = query.shape
+ if head_dim not in _HEAD_DIMS or key.shape[-1] != head_dim or value.shape != key.shape:
+ return "head_dim"
+ kv_heads, key_len = key.shape[1], key.shape[2]
+ if key.shape[0] != batch:
+ return "shape"
+ if kv_heads != heads and not (enable_gqa and heads % kv_heads == 0):
+ return "heads"
+ if query_len < _MIN_QUERY_LEN or key_len < _MIN_KEY_LEN:
+ return "short"
+ if query.stride(-1) != 1 or key.stride(-1) != 1 or value.stride(-1) != 1:
+ return "layout"
+ if torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad):
+ return "grad"
+ if torch.is_autocast_enabled("cuda"):
+ return "autocast"
+ return None
+
+
+def _check_first_result(
+ index: int,
+ original: Callable[..., torch.Tensor],
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ scale: float | None,
+ out: torch.Tensor,
+) -> None:
+ """Raise unless the first SageAttention result of this kernel variant agrees with PyTorch's.
+
+ SageAttention launches its CUDA kernels without checking for errors, so a kernel the installed build lacks for
+ this GPU raises nothing: it returns uninitialized memory, and the launch error surfaces in some later, unrelated
+ torch operation -- outside the fallback. Recomputing a slice of the first result with PyTorch, and reading the
+ comparison back to the host, brings both into the caller's ``try``. Once per device, precision and head size --
+ the inputs that pick a different compiled kernel; later calls are not compared. The slice is the first sample's
+ first K/V head group and its last ``_FIRST_CALL_ROWS`` query rows.
+ """
+ group = query.shape[1] // key.shape[1]
+ rows = slice(-_FIRST_CALL_ROWS, None)
+ expected = original(query[:1, :group, rows], key[:1, :1], value[:1, :1], scale=scale, enable_gqa=group > 1)
+ difference = torch.linalg.vector_norm(out[:1, :group, rows] - expected, dtype=torch.float32)
+ reference = torch.linalg.vector_norm(expected, dtype=torch.float32).clamp_min(1e-6)
+ error = (difference / reference).item()
+ if not math.isfinite(error) or error > _FIRST_CALL_MAX_ERROR:
+ raise RuntimeError(f"its first result differs from PyTorch SDPA by {error:.2f} (relative L2)")
+ with _lock:
+ _validated.add((index, query.dtype, query.shape[-1]))
+
+
+def _cache_failure(index: int) -> bool:
+ """Count a Triton cache failure on this device; True once it failed too often in a row to keep retrying."""
+ with _lock:
+ _cache_failures[index] += 1
+ return _cache_failures[index] >= _MAX_CACHE_FAILURES
+
+
+def _disable(index: int, error: Exception, query: torch.Tensor, key: torch.Tensor) -> None:
+ with _lock:
+ first = index not in _disabled_devices
+ _disabled_devices.add(index)
+ if first:
+ logger.warning(
+ f"SageAttention failed on cuda:{index} for query {tuple(query.shape)}, key {tuple(key.shape)}, "
+ f"{query.dtype} ({error!r}). Attention on this device uses PyTorch SDPA until restart; set "
+ "`attention_backend: auto` to stop using SageAttention."
+ )
+
+
+def _announce(index: int, version: str, kernel: Callable[..., torch.Tensor]) -> None:
+ if index in _announced_devices:
+ return
+ with _lock:
+ if index in _announced_devices:
+ return
+ _announced_devices.add(index)
+ major, minor = torch.cuda.get_device_capability(index)
+ name = getattr(kernel, "func", kernel).__name__
+ logger.info(
+ f"SageAttention {version} serves diffusion-model attention on cuda:{index} "
+ f"({torch.cuda.get_device_name(index)}, sm_{major}{minor}, {name})."
+ )
+
+
+def install_sage_attention() -> None:
+ """Rebind ``F.scaled_dot_product_attention`` so ``sage_attention_scope`` can route calls to SageAttention.
+
+ Idempotent and never raises: a machine that cannot run SageAttention logs why and keeps PyTorch's SDPA. Never
+ installs on ROCm, which also keeps ``rocm_sdpa_chunks_math`` reading the ROCm guard's sentinel on the live binding.
+ Installed once at startup via ``apply_monkeypatches`` when ``attention_backend`` is ``sage``.
+ """
+ from importlib.metadata import PackageNotFoundError
+
+ functional = torch.nn.functional
+ if getattr(functional.scaled_dot_product_attention, _SENTINEL, False):
+ return
+ unavailable = "attention_backend is 'sage', but SageAttention is unavailable here: "
+ see_docs = f"See {DOCS_URL}. Using PyTorch SDPA."
+ try:
+ if torch.version.hip is not None or torch.version.cuda is None or not torch.cuda.is_available():
+ logger.warning(unavailable + "it needs an NVIDIA GPU with CUDA. Using PyTorch SDPA.")
+ return
+ capabilities = {index: torch.cuda.get_device_capability(index) for index in range(torch.cuda.device_count())}
+ if not any(capability >= (8, 0) for capability in capabilities.values()):
+ logger.warning(
+ unavailable + "it needs an NVIDIA GPU with compute capability 8.0 or newer. Using PyTorch SDPA."
+ )
+ return
+ version, sageattention = _import_sageattention()
+ parsed = _parse_version(version)
+ if parsed is None or parsed < _MIN_VERSION:
+ logger.warning(
+ unavailable + f"version {version} is installed, and {'.'.join(map(str, _MIN_VERSION))} or newer is "
+ f"required (the `sageattention` release on PyPI is SageAttention 1). {see_docs}"
+ )
+ return
+ cuda = _parse_version(torch.version.cuda) or (0, 0, 0)
+ kernels = {
+ index: kernel for index, cap in capabilities.items() if (kernel := _kernel(sageattention, cap, cuda[:2]))
+ }
+ for kernel in kernels.values():
+ if unaccepted := _unaccepted_arguments(kernel):
+ logger.warning(
+ unavailable + f"its {kernel.func.__name__} does not take {', '.join(unaccepted)}, so this "
+ f"SageAttention {version} differs from the version InvokeAI was validated with. {see_docs}"
+ )
+ return
+ except PackageNotFoundError:
+ logger.warning(
+ unavailable + f"`sageattention` imports, but has no package metadata to check its version. {see_docs}"
+ )
+ return
+ except ModuleNotFoundError as e:
+ if (e.name or "").split(".")[0] == "triton":
+ hint = "it needs Triton, which on Windows is the `triton-windows` package"
+ else:
+ hint = "the `sageattention` package is not installed"
+ logger.warning(unavailable + f"{hint}. {see_docs}")
+ return
+ except Exception as e:
+ logger.warning(unavailable + f"loading it failed ({e!r}). {see_docs}")
+ return
+
+ original = functional.scaled_dot_product_attention
+
+ @functools.wraps(original)
+ def sage_scaled_dot_product_attention(
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_mask: torch.Tensor | None = None,
+ dropout_p: float = 0.0,
+ is_causal: bool = False,
+ *args: Any,
+ scale: float | None = None,
+ enable_gqa: bool = False,
+ **kwargs: Any,
+ ) -> torch.Tensor:
+ usage = _scope.get()
+ if usage is not None:
+ reason = _ineligible(query, key, value, attn_mask, dropout_p, is_causal, enable_gqa, bool(args or kwargs))
+ if reason is None and usage.memory_short:
+ reason = "memory"
+ if reason is None:
+ index = _cuda_index(query, key, value)
+ kernel = kernels.get(index) if index is not None and index not in _disabled_devices else None
+ if index is None or kernel is None:
+ reason = "device"
+ else:
+ try:
+ with torch.cuda.device(index):
+ out = kernel(
+ query,
+ key,
+ value,
+ tensor_layout="HND",
+ is_causal=False,
+ sm_scale=scale if scale is not None else query.shape[-1] ** -0.5,
+ )
+ if (index, query.dtype, query.shape[-1]) not in _validated:
+ _check_first_result(index, original, query, key, value, scale, out)
+ except Exception as e:
+ if isinstance(e, RuntimeError) and is_oom_error(e):
+ # SageAttention's quantized copies did not fit. SDPA needs less; the rest of this scope
+ # uses it rather than flushing the allocator's cache on every call, and the next
+ # generation tries SageAttention again. A CUDA fault that poisons the context is no OOM:
+ # it retires the device below, and SDPA then fails the same way.
+ usage.memory_short = True
+ reason = "memory"
+ logger.info(
+ f"SageAttention ran out of memory on cuda:{index} for query {tuple(query.shape)}; "
+ "the rest of this generation uses PyTorch SDPA."
+ )
+ elif isinstance(e, OSError) and not _cache_failure(index):
+ reason = "cache"
+ else:
+ _disable(index, e, query, key)
+ reason = "error"
+ else:
+ if _cache_failures[index]:
+ with _lock:
+ _cache_failures.pop(index, None)
+ usage.served += 1
+ _announce(index, version, kernel)
+ return out
+ usage.fell_back[reason] += 1
+ return original(
+ query, key, value, attn_mask, dropout_p, is_causal, *args, scale=scale, enable_gqa=enable_gqa, **kwargs
+ )
+
+ setattr(sage_scaled_dot_product_attention, _SENTINEL, True)
+ functional.scaled_dot_product_attention = sage_scaled_dot_product_attention
+ unserved = sorted(index for index in capabilities if index not in kernels)
+ logger.info(
+ f"SageAttention {version} enabled for the attention inside supported diffusion models "
+ "(attention_backend: sage). Text encoders, VAEs and masked attention keep PyTorch SDPA"
+ + (f"; no SageAttention kernel for cuda:{', cuda:'.join(map(str, unserved))}." if unserved else ".")
+ )
diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json
index 06c65b30c3d..357abb70503 100644
--- a/invokeai/frontend/api/openapi.json
+++ b/invokeai/frontend/api/openapi.json
@@ -59649,6 +59649,13 @@
"description": "Slice size, valid when attention_type==\"sliced\".",
"default": "auto"
},
+ "attention_backend": {
+ "type": "string",
+ "enum": ["auto", "sage"],
+ "title": "Attention Backend",
+ "description": "Attention kernel for the attention inside diffusion models. `auto` lets PyTorch choose among its exact SDPA kernels and never selects a quantized one. `sage` uses SageAttention 2 for the models listed in the SageAttention docs: faster, most at high resolutions and for video, but quantized, so images differ slightly from `auto` at the same seed. SageAttention is installed separately and needs an NVIDIA GPU with compute capability 8.0 or newer; masked attention, text encoders and VAEs keep PyTorch SDPA.",
+ "default": "auto"
+ },
"force_tiled_decode": {
"type": "boolean",
"title": "Force Tiled Decode",
@@ -59979,7 +59986,7 @@
"additionalProperties": false,
"type": "object",
"title": "InvokeAIAppConfig",
- "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n show_donation_link: Show the Donate to InvokeAI link in the frontend menus.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n db_url: URL of a MySQL 8.4+ or MariaDB 10.11+ database to use instead of the SQLite database in `db_dir`, e.g. `mariadb+pymysql://invokeai:password@db.example/invokeai`. Needs the `mysql` extra. One InvokeAI process uses a database at a time. Read at startup only.\n db_synchronous: SQLite durability setting. `full`, the default and what InvokeAI has always used, flushes every commit to disk. `normal` acknowledges commits without waiting for that flush - measured at roughly 12x shorter commits on an SSD - and cannot corrupt the database under WAL, which is why it is refused, with a warning, when WAL is unavailable for the database file. What `normal` gives up is the most recent transactions on a power loss or OS crash: a just-written image record or queue status, not the image file itself.
Valid values: `full`, `normal`\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n reserve_vram_gb: The amount of VRAM (in GB) to subtract from the model cache's available memory budget. Defaults to 0.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n offload_on_architecture_switch: When a session uses main models of another architecture than the previous one on the same GPU, move every model it does not use from VRAM to RAM before it runs. Models both sessions use stay. Switching back reloads the moved models from RAM, which costs a few seconds per switch when both architectures would have fit in VRAM together; turn this off if you alternate between such models.\n clear_vram_after_session: After every queue item (each image of a batch), move the models held in VRAM to RAM and return the freed memory to the GPU driver. The models stay cached in RAM; the next item reloads them from there, which costs a few seconds (more with keep_ram_copy_of_weights off). Changes nothing where VRAM is system memory (CPU, MPS, integrated GPUs). Turn this on if a generation that runs fine the first time runs out of memory, or slows down sharply because the GPU spills into shared memory, when it is repeated.\n fp8_compute: Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU, or a ROCm GPU that passes an fp8 matmul probe; falls back automatically otherwise). Speeds up denoising and keeps the transformer at fp8 size - FP8 Storage alone also keeps it at that size, without the speedup - but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM).\n fp8_compute_full_precision_hints: Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, \"expandable_segments:True\" is used.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number)\n precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n rocm_aotriton_experimental: Use AOTriton's fused (flash and memory-efficient) attention kernels on AMD GPUs that PyTorch marks experimental for them, such as the RX 7600/7700/7800 series, Ryzen AI 300/Max and the RX 9060 series, by setting TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL. Without them, attention on those GPUs runs on the slower math kernel, which also needs far more memory. `auto` turns them on only with a ROCm 10 build of PyTorch, and only when every GPU used for generation is one they were measured correct on (so far gfx1200, the RX 9060 series); `on` and `off` decide for any build and GPU. A TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL value already set in the environment takes precedence. Has no effect on other GPUs.
Valid values: `auto`, `on`, `off`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n noise_dtype: The dtype seeded noise is drawn in for SD1.5/SDXL, SD3, FLUX.1, FLUX.2, CogView4 and the Z-Image seed variance enhancer. `float32` draws the same noise on every platform. `float16` is the half-precision draw of earlier versions: on macOS it keeps the images your seeds gave before this update; on Windows and Linux both values give (nearly) the same images, and images made before the update cannot be reproduced there.
Valid values: `float32`, `float16`\n wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation.\n pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.\n auto_tiled_decode: Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n allow_private_download_urls: Allow the download queue to fetch from loopback, link-local and private-network addresses. Disabled by default so that a download URL cannot be used to reach services that are only reachable from the server. Enable this only if you install models from a mirror on your own network.\n download_proxy: Optional HTTP proxy for model downloads. The proxy must enforce the public-address policy because proxy-side DNS cannot be checked by InvokeAI.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n image_index_enabled: Maintain a semantic embedding index of gallery images and videos, used by the image map and semantic search features.\n image_index_model: Name of the installed CLIP Vision or SigLIP model used to embed gallery images and video thumbnails. Changing the model discards embeddings computed by the previous model.\n image_index_device: Set to `cpu` to compute image embeddings on the CPU with a service-local copy of the model - avoids VRAM use and lets indexing run during generations. Any other value is ignored: embeddings otherwise run on the model cache's device, pausing while generations are in progress.\n image_index_batch_size: Number of gallery items embedded per batch by the image index worker.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_atlascloud_api_key: API key for Atlas Cloud image generation.\n external_atlascloud_base_url: Base URL override for Atlas Cloud image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set.\n http_compression_level: Compression level for gzipped HTTP API responses. 0 disables response compression entirely, 1 is fastest, 9 (the default) is smallest. Compression runs on the event loop and blocks the whole server while it works, and level 9 costs about 5.5x the time of level 1 for 0.4 percentage points of extra compression, so lowering this makes the app noticeably more responsive on large libraries. Set to 0 when a reverse proxy already compresses responses."
+ "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n show_donation_link: Show the Donate to InvokeAI link in the frontend menus.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n db_url: URL of a MySQL 8.4+ or MariaDB 10.11+ database to use instead of the SQLite database in `db_dir`, e.g. `mariadb+pymysql://invokeai:password@db.example/invokeai`. Needs the `mysql` extra. One InvokeAI process uses a database at a time. Read at startup only.\n db_synchronous: SQLite durability setting. `full`, the default and what InvokeAI has always used, flushes every commit to disk. `normal` acknowledges commits without waiting for that flush - measured at roughly 12x shorter commits on an SSD - and cannot corrupt the database under WAL, which is why it is refused, with a warning, when WAL is unavailable for the database file. What `normal` gives up is the most recent transactions on a power loss or OS crash: a just-written image record or queue status, not the image file itself.
Valid values: `full`, `normal`\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n reserve_vram_gb: The amount of VRAM (in GB) to subtract from the model cache's available memory budget. Defaults to 0.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n offload_on_architecture_switch: When a session uses main models of another architecture than the previous one on the same GPU, move every model it does not use from VRAM to RAM before it runs. Models both sessions use stay. Switching back reloads the moved models from RAM, which costs a few seconds per switch when both architectures would have fit in VRAM together; turn this off if you alternate between such models.\n clear_vram_after_session: After every queue item (each image of a batch), move the models held in VRAM to RAM and return the freed memory to the GPU driver. The models stay cached in RAM; the next item reloads them from there, which costs a few seconds (more with keep_ram_copy_of_weights off). Changes nothing where VRAM is system memory (CPU, MPS, integrated GPUs). Turn this on if a generation that runs fine the first time runs out of memory, or slows down sharply because the GPU spills into shared memory, when it is repeated.\n fp8_compute: Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU, or a ROCm GPU that passes an fp8 matmul probe; falls back automatically otherwise). Speeds up denoising and keeps the transformer at fp8 size - FP8 Storage alone also keeps it at that size, without the speedup - but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM).\n fp8_compute_full_precision_hints: Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally. Unset on a ROCm build of PyTorch on Windows, \"expandable_segments:True\" is used.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number)\n precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n rocm_aotriton_experimental: Use AOTriton's fused (flash and memory-efficient) attention kernels on AMD GPUs that PyTorch marks experimental for them, such as the RX 7600/7700/7800 series, Ryzen AI 300/Max and the RX 9060 series, by setting TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL. Without them, attention on those GPUs runs on the slower math kernel, which also needs far more memory. `auto` turns them on only with a ROCm 10 build of PyTorch, and only when every GPU used for generation is one they were measured correct on (so far gfx1200, the RX 9060 series); `on` and `off` decide for any build and GPU. A TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL value already set in the environment takes precedence. Has no effect on other GPUs.
Valid values: `auto`, `on`, `off`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n noise_dtype: The dtype seeded noise is drawn in for SD1.5/SDXL, SD3, FLUX.1, FLUX.2, CogView4 and the Z-Image seed variance enhancer. `float32` draws the same noise on every platform. `float16` is the half-precision draw of earlier versions: on macOS it keeps the images your seeds gave before this update; on Windows and Linux both values give (nearly) the same images, and images made before the update cannot be reproduced there.
Valid values: `float32`, `float16`\n wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation.\n pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n attention_backend: Attention kernel for the attention inside diffusion models. `auto` lets PyTorch choose among its exact SDPA kernels and never selects a quantized one. `sage` uses SageAttention 2 for the models listed in the SageAttention docs: faster, most at high resolutions and for video, but quantized, so images differ slightly from `auto` at the same seed. SageAttention is installed separately and needs an NVIDIA GPU with compute capability 8.0 or newer; masked attention, text encoders and VAEs keep PyTorch SDPA.
Valid values: `auto`, `sage`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.\n auto_tiled_decode: Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n allow_private_download_urls: Allow the download queue to fetch from loopback, link-local and private-network addresses. Disabled by default so that a download URL cannot be used to reach services that are only reachable from the server. Enable this only if you install models from a mirror on your own network.\n download_proxy: Optional HTTP proxy for model downloads. The proxy must enforce the public-address policy because proxy-side DNS cannot be checked by InvokeAI.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n image_index_enabled: Maintain a semantic embedding index of gallery images and videos, used by the image map and semantic search features.\n image_index_model: Name of the installed CLIP Vision or SigLIP model used to embed gallery images and video thumbnails. Changing the model discards embeddings computed by the previous model.\n image_index_device: Set to `cpu` to compute image embeddings on the CPU with a service-local copy of the model - avoids VRAM use and lets indexing run during generations. Any other value is ignored: embeddings otherwise run on the model cache's device, pausing while generations are in progress.\n image_index_batch_size: Number of gallery items embedded per batch by the image index worker.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_atlascloud_api_key: API key for Atlas Cloud image generation.\n external_atlascloud_base_url: Base URL override for Atlas Cloud image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set.\n http_compression_level: Compression level for gzipped HTTP API responses. 0 disables response compression entirely, 1 is fastest, 9 (the default) is smallest. Compression runs on the event loop and blocks the whole server while it works, and level 9 costs about 5.5x the time of level 1 for 0.4 percentage points of extra compression, so lowering this makes the app noticeably more responsive on large libraries. Set to 0 when a reverse proxy already compresses responses."
},
"InvokeAIAppConfigWithSetFields": {
"properties": {
diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts
index f4c062c0dbc..755a2516c95 100644
--- a/invokeai/frontend/api/schema.ts
+++ b/invokeai/frontend/api/schema.ts
@@ -22672,6 +22672,7 @@ export type components = {
* pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.
* attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`
* attention_slice_size: Slice size, valid when attention_type=="sliced".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`
+ * attention_backend: Attention kernel for the attention inside diffusion models. `auto` lets PyTorch choose among its exact SDPA kernels and never selects a quantized one. `sage` uses SageAttention 2 for the models listed in the SageAttention docs: faster, most at high resolutions and for video, but quantized, so images differ slightly from `auto` at the same seed. SageAttention is installed separately and needs an NVIDIA GPU with compute capability 8.0 or newer; masked attention, text encoders and VAEs keep PyTorch SDPA.
Valid values: `auto`, `sage`
* force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.
* auto_tiled_decode: Decode large images in tiles when an untiled decode's estimated working memory would take most of the GPU's memory (FLUX.1, Z-Image, Qwen-Image, Krea-2 and Wan decodes). Turn off to decode untiled unless tiling is requested; the FLUX.1 and Z-Image decodes then still retry tiled after running out of memory. Anima keeps tiling a decode too large for its GPU either way, because tiling is what makes that case fast there.
* pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.
@@ -23141,6 +23142,13 @@ export type components = {
* @enum {unknown}
*/
attention_slice_size?: "auto" | "balanced" | "max" | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8;
+ /**
+ * Attention Backend
+ * @description Attention kernel for the attention inside diffusion models. `auto` lets PyTorch choose among its exact SDPA kernels and never selects a quantized one. `sage` uses SageAttention 2 for the models listed in the SageAttention docs: faster, most at high resolutions and for video, but quantized, so images differ slightly from `auto` at the same seed. SageAttention is installed separately and needs an NVIDIA GPU with compute capability 8.0 or newer; masked attention, text encoders and VAEs keep PyTorch SDPA.
+ * @default auto
+ * @enum {string}
+ */
+ attention_backend?: "auto" | "sage";
/**
* Force Tiled Decode
* @description Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty). A tiled decode is not pixel-identical to a single-pass one: a VAE decoder normalises and attends over the whole image, so the difference is spread across it rather than confined to the tile seams. Also tiles the encode. Read by the SD1.5/SDXL, FLUX.1, Z-Image and Qwen-Image VAE nodes, the Anima encode and the Krea-2 style reference; other model families ignore it.
diff --git a/tests/backend/util/data/qwen_image_first_block_outlier_head.safetensors b/tests/backend/util/data/qwen_image_first_block_outlier_head.safetensors
new file mode 100644
index 00000000000..2306ea05e6e
Binary files /dev/null and b/tests/backend/util/data/qwen_image_first_block_outlier_head.safetensors differ
diff --git a/tests/backend/util/test_sage_attention.py b/tests/backend/util/test_sage_attention.py
new file mode 100644
index 00000000000..ca1ae56a509
--- /dev/null
+++ b/tests/backend/util/test_sage_attention.py
@@ -0,0 +1,788 @@
+"""SageAttention routing for `F.scaled_dot_product_attention` (`backend.util.sage_attention`).
+
+The wrapper must hand exactly the eligible calls made inside a `sage_attention_scope` to the SageAttention kernel for
+the call's GPU, pass every other call to torch with its arguments untouched, keep one thread's scope from leaking into
+another, check each GPU's first result against torch, fall back for good on a device whose kernel failed, and install
+only where SageAttention can run at all.
+
+CPU tensors stand in for the CUDA device and a fake `sageattention` module for the kernels, so none of this establishes
+how the real kernels behave; `TestOnNvidiaHardware` does, on a GPU with SageAttention installed (`-m slow`).
+"""
+
+import contextlib
+import functools
+import importlib.util
+import logging
+import threading
+from collections import Counter
+from importlib.metadata import PackageNotFoundError
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+import invokeai.backend.util.sage_attention as sage
+from invokeai.backend.util.sage_attention import install_sage_attention, sage_attention_scope
+
+# torch's own function, captured before any test installs the wrapper.
+TORCH_SDPA = F.scaled_dot_product_attention
+VERSION = "2.2.0+cu130torch2.10.0andhigher.post6"
+KERNELS = (
+ "sageattn_qk_int8_pv_fp16_cuda",
+ "sageattn_qk_int8_pv_fp16_triton",
+ "sageattn_qk_int8_pv_fp8_cuda",
+ "sageattn_qk_int8_pv_fp8_cuda_sm90",
+)
+FIXTURE = Path(__file__).parent / "data" / "qwen_image_first_block_outlier_head.safetensors"
+
+
+def _reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, scale: float | None = None) -> torch.Tensor:
+ """Attention in float32 over [B, H, S, D], K/V heads repeated for GQA."""
+ q, k, v = q.float(), k.float(), v.float()
+ if k.shape[1] != q.shape[1]:
+ k, v = (t.repeat_interleave(q.shape[1] // t.shape[1], dim=1) for t in (k, v))
+ scale = q.shape[-1] ** -0.5 if scale is None else scale
+ return torch.softmax(q @ k.transpose(-1, -2) * scale, dim=-1) @ v
+
+
+class FakeSageAttention:
+ """Stands in for the `sageattention` module. Its kernels take the real kernels' named parameters, record each
+ call and return the exact result -- or `result`, when a test sets one."""
+
+ def __init__(self, without: tuple[str, ...] = ()) -> None:
+ self.calls: list[dict] = []
+ self.error: Exception | None = None
+ self.result: torch.Tensor | None = None
+ for name in KERNELS:
+ if name not in without:
+ setattr(self, name, self._kernel(name))
+
+ def _kernel(self, name: str):
+ def kernel(
+ q, k, v, tensor_layout="HND", is_causal=False, qk_quant_gran="per_thread", sm_scale=None,
+ pv_accum_dtype="fp32", smooth_k=True, smooth_v=False, return_lse=False, **kwargs,
+ ): # fmt: skip
+ self.calls.append(
+ {"kernel": name, "layout": tensor_layout, "causal": is_causal, "sm_scale": sm_scale, "k": k.shape,
+ "qk_quant_gran": qk_quant_gran, "pv_accum_dtype": pv_accum_dtype, "smooth_k": smooth_k}
+ ) # fmt: skip
+ if self.error is not None:
+ raise self.error
+ if self.result is not None:
+ return self.result
+ return _reference(q, k, v, sm_scale).to(q.dtype)
+
+ kernel.__name__ = name
+ return kernel
+
+
+def _q(*shape: int, dtype: torch.dtype = torch.float16) -> torch.Tensor:
+ return torch.randn(*shape).to(dtype)
+
+
+# Long enough for SageAttention by every rule, small enough to be cheap on the CPU.
+QUERY_LEN = sage._MIN_QUERY_LEN
+KEY_LEN = sage._MIN_KEY_LEN
+
+
+def _cuda_build(
+ monkeypatch, capabilities: dict[int, tuple[int, int]], module: FakeSageAttention | None = None
+) -> tuple[FakeSageAttention, list[dict]]:
+ """Make this build look like CUDA 13 with the given devices and the fake module the installed SageAttention, with
+ every device's first-call check already passed. Returns the module and the argument tuples of every call that
+ reached torch's function (computed in float32: no CPU half-precision GEMM)."""
+ module = module or FakeSageAttention()
+ torch_calls: list[dict] = []
+
+ def torch_spy(
+ query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *args, scale=None, enable_gqa=False, **kwargs
+ ):
+ torch_calls.append(
+ {"mask": attn_mask, "dropout": dropout_p, "causal": is_causal, "scale": scale, "gqa": enable_gqa,
+ "args": args, "kwargs": kwargs, "q": query.shape}
+ ) # fmt: skip
+ mask = attn_mask if attn_mask is None or attn_mask.dtype == torch.bool else attn_mask.float()
+ out = TORCH_SDPA(
+ query.float(), key.float(), value.float(), mask, dropout_p, is_causal, scale=scale, enable_gqa=enable_gqa
+ )
+ return out.to(query.dtype)
+
+ monkeypatch.setattr(F, "scaled_dot_product_attention", torch_spy)
+ monkeypatch.setattr(torch.version, "hip", None)
+ monkeypatch.setattr(torch.version, "cuda", "13.0")
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(torch.cuda, "device_count", lambda: len(capabilities))
+ monkeypatch.setattr(torch.cuda, "get_device_capability", lambda index=None: capabilities[index or 0])
+ monkeypatch.setattr(torch.cuda, "get_device_name", lambda index=None: "Fake GPU")
+ monkeypatch.setattr(sage, "_import_sageattention", lambda: (VERSION, module))
+ monkeypatch.setattr(sage, "_cuda_index", lambda q, k, v: 0)
+ monkeypatch.setattr(sage, "_disabled_devices", set())
+ monkeypatch.setattr(sage, "_validated", _all_variants(capabilities))
+ monkeypatch.setattr(sage, "_cache_failures", Counter())
+ monkeypatch.setattr(sage, "_announced_devices", set())
+ monkeypatch.setattr(torch.cuda, "device", lambda index: contextlib.nullcontext())
+ return module, torch_calls
+
+
+def _all_variants(devices) -> set:
+ """Every kernel variant of these devices, as if each had passed its first-call check."""
+ return {
+ (index, dtype, head_dim) for index in devices for dtype in sage._HALF_PRECISION for head_dim in sage._HEAD_DIMS
+ }
+
+
+@pytest.fixture
+def installed(monkeypatch):
+ """The wrapper installed on a build that looks like CUDA 13 with one sm_89 device."""
+ module, torch_calls = _cuda_build(monkeypatch, {0: (8, 9)})
+ spy = F.scaled_dot_product_attention
+ install_sage_attention()
+ assert F.scaled_dot_product_attention is not spy
+ return module, torch_calls
+
+
+class TestKernelTable:
+ @pytest.mark.parametrize(
+ ("capability", "cuda", "kernel", "options"),
+ [
+ ((8, 0), (13, 0), "sageattn_qk_int8_pv_fp16_cuda", {"pv_accum_dtype": "fp32"}),
+ # Upstream's `sageattn` picks the Triton kernel here; the Windows build moved sm86 to CUDA unexplained.
+ ((8, 6), (13, 0), "sageattn_qk_int8_pv_fp16_triton", {}),
+ ((8, 9), (13, 0), "sageattn_qk_int8_pv_fp8_cuda", {"pv_accum_dtype": "fp32+fp16"}),
+ ((8, 9), (12, 6), "sageattn_qk_int8_pv_fp8_cuda", {"pv_accum_dtype": "fp32+fp32"}),
+ ((9, 0), (13, 0), "sageattn_qk_int8_pv_fp8_cuda_sm90", {"pv_accum_dtype": "fp32+fp32"}),
+ ((12, 0), (13, 0), "sageattn_qk_int8_pv_fp8_cuda", {"pv_accum_dtype": "fp32+fp16", "qk_quant_gran": "per_warp"}),
+ ((12, 0), (12, 6), "sageattn_qk_int8_pv_fp8_cuda", {"pv_accum_dtype": "fp32", "qk_quant_gran": "per_warp"}),
+ ],
+ ids=["ampere-a100", "ampere-rtx30", "ada", "ada-cuda12.6", "hopper", "blackwell", "blackwell-cuda12.6"],
+ ) # fmt: skip
+ def test_each_generation_gets_its_kernel_with_smoothing_off(self, capability, cuda, kernel, options):
+ """K smoothing broke the first block of Qwen-Image and Krea-2 (cosine 0.18 and 0.53); FP16 accumulation inside
+ the FP8 kernels needs CUDA 12.8. See `_kernel`."""
+ module = FakeSageAttention()
+ q = _q(1, 2, 8, 64)
+
+ sage._kernel(module, capability, cuda)(q, q, q, tensor_layout="HND", is_causal=False, sm_scale=None)
+
+ assert module.calls[0]["kernel"] == kernel
+ assert {name: module.calls[0][name] for name in (*options, "smooth_k")} == options | {"smooth_k": False}
+
+ @pytest.mark.parametrize(
+ "capability", [(7, 5), (8, 7), (10, 3), (11, 0)], ids=["turing", "jetson-orin", "sm103", "sm110"]
+ )
+ def test_gpus_sageattention_has_no_kernel_for_get_none(self, capability):
+ assert sage._kernel(FakeSageAttention(), capability, (13, 0)) is None
+
+
+class TestRouting:
+ def test_an_eligible_call_in_scope_runs_on_the_kernel_in_hnd_layout(self, installed):
+ module, torch_calls = installed
+ q, k, v = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+
+ with sage_attention_scope():
+ out = F.scaled_dot_product_attention(q, k, v)
+
+ assert torch_calls == []
+ call = module.calls[0]
+ assert (call["layout"], call["causal"], call["sm_scale"], call["k"]) == ("HND", False, 64**-0.5, k.shape)
+ torch.testing.assert_close(out.float(), _reference(q, k, v), atol=2e-3, rtol=2e-3)
+
+ def test_an_explicit_scale_reaches_the_kernel(self, installed):
+ module, _ = installed
+ q, k = _q(1, 2, QUERY_LEN, 128), _q(1, 2, KEY_LEN, 128)
+
+ with sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, k, scale=0.3)
+
+ assert module.calls[0]["sm_scale"] == 0.3
+
+ def test_outside_a_scope_every_call_reaches_torch_untouched(self, installed):
+ """Text encoders and VAEs run outside any scope, so they must keep torch's exact semantics."""
+ module, torch_calls = installed
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+
+ F.scaled_dot_product_attention(q, k, k, None, 0.0, False, scale=0.5)
+
+ assert module.calls == []
+ assert torch_calls == [
+ {"mask": None, "dropout": 0.0, "causal": False, "scale": 0.5, "gqa": False, "args": (), "kwargs": {},
+ "q": q.shape}
+ ] # fmt: skip
+
+ @pytest.mark.parametrize(
+ ("make", "reason"),
+ [
+ (lambda: ((_q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)),
+ {"attn_mask": torch.ones(QUERY_LEN, KEY_LEN, dtype=torch.bool)}), "mask"),
+ (lambda: ((_q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)),
+ {"attn_mask": torch.zeros(1, 1, 1, KEY_LEN, dtype=torch.float16)}), "mask"),
+ (lambda: ((_q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)), {"dropout_p": 0.1}), "dropout"),
+ (lambda: ((_q(1, 2, QUERY_LEN, 64), _q(1, 2, QUERY_LEN, 64)), {"is_causal": True}), "causal"),
+ (lambda: ((_q(1, 2, QUERY_LEN, 64, dtype=torch.float32), _q(1, 2, KEY_LEN, 64, dtype=torch.float32)), {}),
+ "dtype"),
+ (lambda: ((_q(1, 2, QUERY_LEN, 40), _q(1, 2, KEY_LEN, 40)), {}), "head_dim"),
+ (lambda: ((_q(1, 2, QUERY_LEN, 80), _q(1, 2, KEY_LEN, 80)), {}), "head_dim"),
+ (lambda: ((_q(1, 2, QUERY_LEN, 256), _q(1, 2, KEY_LEN, 256)), {}), "head_dim"),
+ (lambda: ((_q(2, QUERY_LEN, 64), _q(2, KEY_LEN, 64)), {}), "shape"),
+ (lambda: ((_q(2, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)), {}), "shape"),
+ (lambda: ((_q(1, 2, QUERY_LEN - 1, 64), _q(1, 2, KEY_LEN, 64)), {}), "short"),
+ (lambda: ((_q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN - 1, 64)), {}), "short"),
+ (lambda: ((_q(1, 2, QUERY_LEN, 128)[..., ::2], _q(1, 2, KEY_LEN, 64)), {}), "layout"),
+ ],
+ ids=["bool-mask", "additive-mask", "dropout", "causal", "fp32", "hd40", "hd80", "hd256", "3d",
+ "broadcast-batch", "short-query", "short-key", "strided-last-dim"],
+ ) # fmt: skip
+ def test_ineligible_calls_reach_torch_with_their_arguments(self, installed, caplog, make, reason):
+ module, torch_calls = installed
+ (q, k), kwargs = make()
+ caplog.set_level(logging.DEBUG, logger=sage.__name__)
+
+ with sage_attention_scope():
+ out = F.scaled_dot_product_attention(q, k, k, **kwargs)
+
+ assert module.calls == []
+ assert len(torch_calls) == 1
+ assert torch_calls[0]["mask"] is kwargs.get("attn_mask")
+ assert torch_calls[0]["dropout"] == kwargs.get("dropout_p", 0.0)
+ assert torch_calls[0]["causal"] == kwargs.get("is_causal", False)
+ assert torch_calls[0]["gqa"] == kwargs.get("enable_gqa", False)
+ assert out.dtype == q.dtype
+ assert any(f"on PyTorch SDPA: {reason}=1" in r.message for r in caplog.records)
+
+ @pytest.mark.parametrize(("kv_heads", "enable_gqa"), [(2, False), (3, True)], ids=["gqa-flag-off", "indivisible"])
+ def test_mismatched_heads_get_torchs_error_not_a_kernel_answer(self, installed, kv_heads, enable_gqa):
+ module, torch_calls = installed
+ q, k = _q(1, 8, QUERY_LEN, 64), _q(1, kv_heads, KEY_LEN, 64)
+
+ with sage_attention_scope(), pytest.raises(RuntimeError):
+ F.scaled_dot_product_attention(q, k, k, enable_gqa=enable_gqa)
+
+ assert module.calls == [] and len(torch_calls) == 1
+
+ def test_an_input_that_requires_grad_under_grad_mode_reaches_torch(self, installed):
+ module, torch_calls = installed
+ q = _q(1, 2, QUERY_LEN, 64).requires_grad_()
+ k = _q(1, 2, KEY_LEN, 64)
+
+ with torch.enable_grad(), sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, k)
+
+ assert module.calls == [] and len(torch_calls) == 1
+
+ def test_a_call_under_cuda_autocast_reaches_torch(self, installed, monkeypatch):
+ """Autocast would hand SageAttention whatever dtype it picks; torch's SDPA applies autocast's rules itself."""
+ module, torch_calls = installed
+ monkeypatch.setattr(torch, "is_autocast_enabled", lambda device_type=None: device_type == "cuda")
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+
+ with sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, k)
+
+ assert module.calls == [] and len(torch_calls) == 1
+
+ @pytest.mark.filterwarnings("ignore:The PyTorch API of nested tensors")
+ def test_a_nested_query_is_not_eligible(self):
+ """Nested (jagged) tensors report four dimensions but have no fixed sequence length to quantize over."""
+ q = torch.nested.nested_tensor([_q(2, QUERY_LEN, 64), _q(2, QUERY_LEN + 1, 64)])
+ k = _q(2, 2, KEY_LEN, 64)
+ assert sage._ineligible(q, k, k, None, 0.0, False, False, False) == "shape"
+
+ def test_an_unknown_keyword_reaches_torch_rather_than_being_dropped(self, installed):
+ module, torch_calls = installed
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+
+ with sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, k, some_future_flag=True)
+
+ assert module.calls == [] and torch_calls[0]["kwargs"] == {"some_future_flag": True}
+
+ def test_grouped_query_attention_reaches_the_kernel_unexpanded(self, installed):
+ module, torch_calls = installed
+ q, k, v = _q(1, 8, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+
+ with sage_attention_scope():
+ out = F.scaled_dot_product_attention(q, k, v, enable_gqa=True)
+
+ assert torch_calls == [] and module.calls[0]["k"] == k.shape
+ torch.testing.assert_close(out.float(), _reference(q, k, v), atol=2e-3, rtol=2e-3)
+
+
+class TestFirstCallCheck:
+ def test_a_devices_first_result_is_compared_with_torch_once(self, monkeypatch):
+ module, torch_calls = _cuda_build(monkeypatch, {0: (8, 9)})
+ monkeypatch.setattr(sage, "_validated", set())
+ install_sage_attention()
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+
+ with sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, k)
+ F.scaled_dot_product_attention(q, k, k)
+
+ assert len(module.calls) == 2, "both calls ran on the kernel"
+ assert len(torch_calls) == 1, "only the first was also computed by torch, for the comparison"
+
+ def test_the_check_pairs_a_gqa_query_group_with_its_own_kv_head(self, monkeypatch):
+ """Query heads 0-3 attend to K/V head 0; comparing them with head 1 instead would retire a correct kernel."""
+ module, torch_calls = _cuda_build(monkeypatch, {0: (8, 9)})
+ monkeypatch.setattr(sage, "_validated", set())
+ install_sage_attention()
+ q, k, v = _q(1, 8, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+ v[:, 1] += 3.0 # a K/V head unlike head 0, so a wrong pairing cannot pass
+
+ with sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, v, enable_gqa=True)
+ F.scaled_dot_product_attention(q, k, v, enable_gqa=True)
+
+ assert len(module.calls) == 2, "the check passed, so the device stays on the kernel"
+ assert len(torch_calls) == 1 and torch_calls[0]["gqa"] and torch_calls[0]["q"] == (1, 4, 1024, 64)
+
+ def test_each_precision_and_head_size_is_checked_on_its_first_call(self, monkeypatch):
+ """They select different compiled kernels; one passing says nothing about the others."""
+ module, torch_calls = _cuda_build(monkeypatch, {0: (8, 9)})
+ monkeypatch.setattr(sage, "_validated", set())
+ install_sage_attention()
+ hd64, hd128 = _q(1, 2, QUERY_LEN, 64), _q(1, 2, QUERY_LEN, 128)
+ bf16 = _q(1, 2, QUERY_LEN, 64, dtype=torch.bfloat16)
+
+ with sage_attention_scope():
+ for q in (hd64, hd64, hd128, hd128, bf16, bf16):
+ F.scaled_dot_product_attention(q, q, q)
+
+ assert len(module.calls) == 6
+ assert len(torch_calls) == 3, "one comparison per variant"
+
+ @pytest.mark.parametrize("garbage", [torch.nan, 50.0], ids=["non-finite", "wrong"])
+ def test_a_wrong_first_result_retires_the_device_and_returns_torchs(self, monkeypatch, caplog, garbage):
+ """A kernel the build lacks for this GPU returns uninitialized memory instead of raising."""
+ module, torch_calls = _cuda_build(monkeypatch, {0: (8, 9)})
+ monkeypatch.setattr(sage, "_validated", set())
+ install_sage_attention()
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+ module.result = torch.full((1, 2, QUERY_LEN, 64), garbage, dtype=torch.float16)
+ caplog.set_level(logging.WARNING, logger=sage.__name__)
+
+ with sage_attention_scope():
+ out = F.scaled_dot_product_attention(q, k, k)
+ F.scaled_dot_product_attention(q, k, k)
+
+ torch.testing.assert_close(out.float(), _reference(q, k, k), atol=2e-3, rtol=2e-3)
+ assert len(module.calls) == 1, "a retired device must not be tried again"
+ assert any("first result differs from PyTorch SDPA" in r.message for r in caplog.records)
+
+
+class TestFailures:
+ def test_a_kernel_failure_falls_back_warns_once_and_retires_only_that_device(self, monkeypatch, caplog):
+ module, _ = _cuda_build(monkeypatch, {0: (8, 9), 1: (8, 9)})
+ install_sage_attention()
+ device = {"index": 0}
+ monkeypatch.setattr(sage, "_cuda_index", lambda q, k, v: device["index"])
+ module.error = RuntimeError("kernel not compiled for this architecture")
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+ caplog.set_level(logging.WARNING, logger=sage.__name__)
+
+ with sage_attention_scope():
+ first = F.scaled_dot_product_attention(q, k, k)
+ F.scaled_dot_product_attention(q, k, k)
+
+ assert len(module.calls) == 1, "a retired device must not be tried again"
+ torch.testing.assert_close(first.float(), _reference(q, k, k), atol=2e-3, rtol=2e-3)
+ warnings = [r.message for r in caplog.records if r.levelno == logging.WARNING]
+ assert len(warnings) == 1 and "cuda:0" in warnings[0]
+
+ module.error = None
+ device["index"] = 1
+ with sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, k)
+ assert len(module.calls) == 2, "another device is unaffected"
+
+ @pytest.mark.parametrize(
+ "error",
+ [
+ torch.OutOfMemoryError("CUDA out of memory. Tried to allocate 1.50 GiB"),
+ RuntimeError("Triton Error [CUDA]: out of memory"),
+ ],
+ ids=["torch-oom", "triton-oom"],
+ )
+ def test_running_out_of_memory_moves_the_rest_of_the_scope_to_torch_only(self, installed, caplog, error):
+ """SDPA needs less memory. Retrying every call would flush the allocator's cache each time; retiring the
+ device would cost every later, smaller generation SageAttention."""
+ module, torch_calls = installed
+ module.error = error
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+ caplog.set_level(logging.INFO, logger=sage.__name__)
+
+ with sage_attention_scope():
+ out = F.scaled_dot_product_attention(q, k, k)
+ module.error = None
+ F.scaled_dot_product_attention(q, k, k)
+ assert len(module.calls) == 1 and len(torch_calls) == 2, "after the OOM, this scope stayed on torch"
+ torch.testing.assert_close(out.float(), _reference(q, k, k), atol=2e-3, rtol=2e-3)
+
+ with sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, k)
+ assert len(module.calls) == 2, "the next scope tries SageAttention again"
+ assert any("ran out of memory" in r.message for r in caplog.records if r.levelno == logging.INFO)
+ assert not [r for r in caplog.records if r.levelno >= logging.WARNING]
+
+ def test_a_triton_cache_failure_falls_back_for_that_call_only(self, installed, caplog):
+ """On Windows two GPUs compiling the same kernel can race for its cache file; that says nothing about it."""
+ module, torch_calls = installed
+ module.error = PermissionError("[WinError 5] Access is denied: 'triton\\cache\\tmp.pid_123'")
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+ caplog.set_level(logging.WARNING, logger=sage.__name__)
+
+ with sage_attention_scope():
+ out = F.scaled_dot_product_attention(q, k, k)
+ module.error = None
+ F.scaled_dot_product_attention(q, k, k)
+
+ torch.testing.assert_close(out.float(), _reference(q, k, k), atol=2e-3, rtol=2e-3)
+ assert len(module.calls) == 2 and len(torch_calls) == 1
+ assert not caplog.records
+
+ def test_only_cache_failures_in_a_row_retire_the_device(self, installed, caplog):
+ """A cache that can never be written would recompile on every call; one that failed now and then is fine."""
+ module, torch_calls = installed
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+ caplog.set_level(logging.WARNING, logger=sage.__name__)
+
+ def call(error):
+ module.error = error
+ with sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, k)
+
+ for error in [PermissionError("denied")] * (sage._MAX_CACHE_FAILURES - 1) + [None]:
+ call(error)
+ for _ in range(sage._MAX_CACHE_FAILURES - 1):
+ call(PermissionError("denied"))
+ assert not caplog.records and 0 not in sage._disabled_devices, "the success in between reset the count"
+
+ call(PermissionError("denied"))
+ call(None)
+ assert 0 in sage._disabled_devices
+ assert len(module.calls) == 2 * sage._MAX_CACHE_FAILURES, "the last call no longer reached the kernel"
+ assert any("cuda:0" in r.message for r in caplog.records if r.levelno == logging.WARNING)
+
+ def test_the_scope_closes_when_its_block_raises(self, installed):
+ module, torch_calls = installed
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+
+ with pytest.raises(ValueError), sage_attention_scope():
+ raise ValueError("generation failed")
+
+ F.scaled_dot_product_attention(q, k, k)
+ assert module.calls == [] and len(torch_calls) == 1
+
+
+class TestScopeIsPerThread:
+ def test_only_the_thread_inside_the_scope_reaches_the_kernel(self, installed):
+ """One generation session per GPU runs concurrently; a scope must not leak into the other session."""
+ module, torch_calls = installed
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+ both_inside = threading.Barrier(2)
+ errors: list[BaseException] = []
+
+ def scoped() -> None:
+ try:
+ with sage_attention_scope():
+ both_inside.wait(timeout=30)
+ F.scaled_dot_product_attention(q, k, k)
+ both_inside.wait(timeout=30)
+ except BaseException as e: # noqa: BLE001 - re-raised in the main thread
+ errors.append(e)
+
+ def unscoped() -> None:
+ try:
+ both_inside.wait(timeout=30)
+ F.scaled_dot_product_attention(q, k, k)
+ both_inside.wait(timeout=30)
+ except BaseException as e: # noqa: BLE001
+ errors.append(e)
+
+ threads = [threading.Thread(target=scoped), threading.Thread(target=unscoped)]
+ for t in threads:
+ t.start()
+ for t in threads:
+ t.join(timeout=60)
+
+ assert errors == []
+ assert len(module.calls) == 1 and len(torch_calls) == 1
+
+
+class TestDeviceRule:
+ @staticmethod
+ def _on(index: int) -> SimpleNamespace:
+ return SimpleNamespace(device=torch.device("cuda", index))
+
+ def test_tensors_on_one_cuda_device_give_its_index(self):
+ assert sage._cuda_index(self._on(1), self._on(1), self._on(1)) == 1
+
+ def test_mixed_devices_and_cpu_do_not(self):
+ assert sage._cuda_index(self._on(0), self._on(1), self._on(0)) is None
+ cpu = SimpleNamespace(device=torch.device("cpu"))
+ assert sage._cuda_index(cpu, cpu, cpu) is None
+
+ def test_the_kernel_runs_with_the_tensors_device_current(self, monkeypatch):
+ """SageAttention launches on the thread's current device. A single-device install with `device: cuda:1` never
+ makes that device current, so the wrapper does, for the call only."""
+ module, _ = _cuda_build(monkeypatch, {0: (8, 9), 1: (8, 9)})
+ events: list[str] = []
+ kernel = module.sageattn_qk_int8_pv_fp8_cuda
+
+ @functools.wraps(kernel)
+ def recording_kernel(*args, **kwargs):
+ events.append("kernel")
+ return kernel(*args, **kwargs)
+
+ @contextlib.contextmanager
+ def recording_device(index):
+ events.append(f"enter cuda:{index}")
+ yield
+ events.append("exit")
+
+ module.sageattn_qk_int8_pv_fp8_cuda = recording_kernel
+ monkeypatch.setattr(torch.cuda, "device", recording_device)
+ monkeypatch.setattr(sage, "_cuda_index", lambda q, k, v: 1)
+ install_sage_attention()
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+
+ with sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, k)
+
+ assert events == ["enter cuda:1", "kernel", "exit"]
+
+
+def _raise(error: Exception):
+ def raiser():
+ raise error
+
+ return raiser
+
+
+class TestInstall:
+ def test_install_is_idempotent(self, monkeypatch):
+ _cuda_build(monkeypatch, {0: (8, 9)})
+ spy = F.scaled_dot_product_attention
+ install_sage_attention()
+ wrapped = F.scaled_dot_product_attention
+ install_sage_attention()
+ assert F.scaled_dot_product_attention is wrapped
+ assert wrapped.__wrapped__ is spy
+
+ @pytest.mark.parametrize(
+ ("patch", "message"),
+ [
+ (lambda m: m.setattr(torch.version, "hip", "7.2"), "NVIDIA GPU with CUDA"),
+ (lambda m: m.setattr(torch.cuda, "is_available", lambda: False), "NVIDIA GPU with CUDA"),
+ (lambda m: m.setattr(torch.cuda, "get_device_capability", lambda index=None: (7, 5)), "8.0 or newer"),
+ (lambda m: m.setattr(sage, "_import_sageattention", _raise(ModuleNotFoundError("x", name="sageattention"))),
+ "`sageattention` package is not installed"),
+ (lambda m: m.setattr(sage, "_import_sageattention", _raise(ModuleNotFoundError("x", name="triton.language"))),
+ "triton-windows"),
+ (lambda m: m.setattr(sage, "_import_sageattention", _raise(PackageNotFoundError("sageattention"))),
+ "no package metadata"),
+ (lambda m: m.setattr(sage, "_import_sageattention", _raise(OSError("DLL load failed"))), "DLL load failed"),
+ (lambda m: m.setattr(sage, "_import_sageattention", lambda: ("1.0.6", FakeSageAttention())),
+ "SageAttention 1"),
+ (lambda m: m.setattr(sage, "_import_sageattention",
+ lambda: (VERSION, FakeSageAttention(without=("sageattn_qk_int8_pv_fp8_cuda",)))),
+ "sageattn_qk_int8_pv_fp8_cuda"),
+ ],
+ ids=["rocm", "no-cuda", "turing", "no-package", "no-triton", "no-metadata", "broken-import", "sageattention-1",
+ "missing-kernel"],
+ ) # fmt: skip
+ def test_install_keeps_torchs_function_and_says_why(self, monkeypatch, caplog, patch, message):
+ _cuda_build(monkeypatch, {0: (8, 9)})
+ spy = F.scaled_dot_product_attention
+ patch(monkeypatch)
+ caplog.set_level(logging.WARNING, logger=sage.__name__)
+
+ install_sage_attention()
+
+ assert F.scaled_dot_product_attention is spy
+ assert any(message in r.message for r in caplog.records), [r.message for r in caplog.records]
+
+ def test_a_kernel_that_would_swallow_smooth_k_is_refused(self, monkeypatch, caplog):
+ """Every real kernel ends in **kwargs, so a renamed `smooth_k` would be ignored and smoothing come back."""
+ module = FakeSageAttention()
+
+ def renamed(q, k, v, tensor_layout="HND", is_causal=False, sm_scale=None, pv_accum_dtype="fp32", **kwargs):
+ raise AssertionError("must not be called")
+
+ renamed.__name__ = "sageattn_qk_int8_pv_fp8_cuda"
+ module.sageattn_qk_int8_pv_fp8_cuda = renamed
+ _cuda_build(monkeypatch, {0: (8, 9)}, module)
+ spy = F.scaled_dot_product_attention
+ caplog.set_level(logging.WARNING, logger=sage.__name__)
+
+ install_sage_attention()
+
+ assert F.scaled_dot_product_attention is spy
+ assert any("does not take smooth_k" in r.message for r in caplog.records)
+
+ @pytest.mark.parametrize(("version", "installs"), [("2.2", True), ("3.0", True), ("2.1.1", False)])
+ def test_versions_are_compared_by_their_numbers(self, monkeypatch, version, installs):
+ _cuda_build(monkeypatch, {0: (8, 9)})
+ monkeypatch.setattr(sage, "_import_sageattention", lambda: (version, FakeSageAttention()))
+
+ install_sage_attention()
+
+ assert getattr(F.scaled_dot_product_attention, sage._SENTINEL, False) is installs
+
+ def test_a_device_without_a_kernel_keeps_torch_while_others_serve(self, monkeypatch, caplog):
+ module, _ = _cuda_build(monkeypatch, {0: (8, 9), 1: (7, 5)})
+ caplog.set_level(logging.INFO, logger=sage.__name__)
+ install_sage_attention()
+ device = {"index": 1}
+ monkeypatch.setattr(sage, "_cuda_index", lambda q, k, v: device["index"])
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+
+ with sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, k)
+ device["index"] = 0
+ F.scaled_dot_product_attention(q, k, k)
+
+ assert len(module.calls) == 1, "only the sm_89 device runs SageAttention"
+ assert any("no SageAttention kernel for cuda:1" in r.message for r in caplog.records)
+
+ @pytest.mark.parametrize("backend", ["sage", "auto"])
+ def test_apply_monkeypatches_installs_sageattention_only_when_configured(self, monkeypatch, backend):
+ from invokeai.app.util.startup_utils import apply_monkeypatches
+
+ _cuda_build(monkeypatch, {0: (8, 9)})
+ monkeypatch.setattr(torch.cuda, "empty_cache", torch.cuda.empty_cache)
+
+ apply_monkeypatches(backend)
+
+ assert getattr(F.scaled_dot_product_attention, sage._SENTINEL, False) is (backend == "sage")
+
+
+class TestLogging:
+ def test_first_use_on_a_device_is_announced_once_with_its_kernel(self, installed, caplog):
+ q, k = _q(1, 2, QUERY_LEN, 64), _q(1, 2, KEY_LEN, 64)
+ caplog.set_level(logging.INFO, logger=sage.__name__)
+
+ with sage_attention_scope():
+ F.scaled_dot_product_attention(q, k, k)
+ F.scaled_dot_product_attention(q, k, k)
+
+ announcements = [r.message for r in caplog.records if "serves diffusion-model attention" in r.message]
+ assert announcements == [
+ f"SageAttention {VERSION} serves diffusion-model attention on cuda:0 "
+ "(Fake GPU, sm_89, sageattn_qk_int8_pv_fp8_cuda)."
+ ]
+
+
+_HAVE_SAGE = importlib.util.find_spec("sageattention") is not None
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(not torch.cuda.is_available() or not _HAVE_SAGE, reason="needs an NVIDIA GPU and sageattention")
+class TestOnNvidiaHardware:
+ """Real kernels: needs CUDA, `sageattention` >= 2.2 and compute capability 8.0+, and about 10 GB of free VRAM for
+ the float32 references. Measured on an RTX 4090."""
+
+ @pytest.fixture(params=["own", "sm86-triton"])
+ def kernel(self, request, monkeypatch):
+ """The kernel this GPU gets, or the Triton kernel an sm86 card gets: it runs on any Ampere or newer GPU, so this
+ one stands in for an sm86 card that is not here."""
+ import sageattention
+
+ capability = torch.cuda.get_device_capability(0)
+ if capability < (8, 0):
+ pytest.skip("SageAttention 2 needs compute capability 8.0 or newer")
+ if request.param == "sm86-triton":
+ pick = sage._kernel
+ monkeypatch.setattr(sage, "_kernel", lambda module, _capability, cuda: pick(module, (8, 6), cuda))
+ monkeypatch.setattr(F, "scaled_dot_product_attention", TORCH_SDPA)
+ monkeypatch.setattr(sage, "_disabled_devices", set())
+ monkeypatch.setattr(sage, "_validated", set())
+ monkeypatch.setattr(sage, "_cache_failures", Counter())
+ monkeypatch.setattr(sage, "_announced_devices", set())
+ install_sage_attention()
+ if not getattr(F.scaled_dot_product_attention, sage._SENTINEL, False):
+ pytest.skip("the installed sageattention is not one InvokeAI uses (see the startup warning)")
+ cuda = sage._parse_version(torch.version.cuda)
+ assert cuda is not None
+ return sage._kernel(sageattention, capability, cuda[:2])
+
+ def test_every_kernel_takes_the_arguments_it_is_given(self, kernel, monkeypatch):
+ import sageattention
+
+ monkeypatch.undo() # the real table, not the fixture's stand-in
+ for capability in ((8, 0), (8, 6), (8, 9), (9, 0), (12, 0)):
+ assert sage._unaccepted_arguments(sage._kernel(sageattention, capability, (13, 0))) == []
+
+ @pytest.mark.parametrize(
+ ("shape_q", "shape_kv", "dtype", "gqa"),
+ [
+ ((1, 24, 4608, 128), (1, 24, 4608, 128), torch.bfloat16, False), # FLUX.1 1024^2
+ ((2, 10, 4096, 64), (2, 10, 4096, 64), torch.float16, False), # SDXL 1024^2, level 1
+ ((1, 48, 4608, 128), (1, 12, 4608, 128), torch.bfloat16, True), # Krea-2 GQA
+ # Lengths that are no multiple of a tile, and cross-attention to a shorter key.
+ ((1, 24, 4429, 128), (1, 24, 845, 128), torch.bfloat16, False),
+ ],
+ ids=["flux", "sdxl", "krea2-gqa", "partial-tile-cross"],
+ )
+ def test_output_stays_close_to_a_float32_reference(self, kernel, shape_q, shape_kv, dtype, gqa):
+ g = torch.Generator(device="cuda").manual_seed(0)
+
+ def make(b, h, s, d): # real activations arrive as [B, S, H, D] storage viewed as [B, H, S, D]
+ return torch.randn(b, s, h, d, device="cuda", generator=g).to(dtype).transpose(1, 2)
+
+ q, k, v = make(*shape_q), make(*shape_kv), make(*shape_kv)
+ with sage_attention_scope():
+ out = F.scaled_dot_product_attention(q, k, v, enable_gqa=gqa)
+
+ # SDPA would pass the bounds below too; the kernel must have served the call.
+ assert torch.equal(out, kernel(q, k, v, tensor_layout="HND", is_causal=False, sm_scale=shape_q[-1] ** -0.5))
+ ref = _reference(q, k, v)
+ assert torch.isfinite(out).all()
+ cos = F.cosine_similarity(out.float().flatten(2), ref.flatten(2), dim=-1)
+ # Measured on random inputs: worst head 0.9991-0.9993; SDPA in bf16 reaches 0.99999.
+ assert cos.min() > 0.998
+ assert (out.float() - ref).norm() / ref.norm() < 0.06
+
+ def test_the_outlier_channel_of_qwen_images_first_block_stays_accurate(self, kernel):
+ """Real activations (see the fixture's metadata). With K smoothing on, as `sageattn` runs it, the relative
+ error here is 1.62; with it off, 0.10."""
+ from safetensors.torch import load_file
+
+ t = {name: tensor.cuda() for name, tensor in load_file(FIXTURE).items()}
+
+ out = kernel(t["q"], t["k"], t["v"], tensor_layout="HND", is_causal=False, sm_scale=128**-0.5)
+
+ ref = _reference(t["q"], t["k"], t["v"])
+ assert (out.float() - ref).norm() / ref.norm() < 0.2
+
+ def test_in_scope_it_is_the_kernel_exactly_and_outside_it_is_torch_exactly(self, kernel):
+ q, k, v = (torch.randn(1, 4608, 24, 128, device="cuda", dtype=torch.bfloat16).transpose(1, 2) for _ in range(3))
+ mask = torch.zeros(1, 1, 4608, 4608, device="cuda", dtype=torch.bfloat16)
+ direct = kernel(q, k, v, tensor_layout="HND", is_causal=False, sm_scale=128**-0.5)
+
+ with sage_attention_scope():
+ assert torch.equal(F.scaled_dot_product_attention(q, k, v), direct)
+ masked = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+ assert torch.equal(masked, TORCH_SDPA(q, k, v, attn_mask=mask))
+ with torch.autocast("cuda", dtype=torch.bfloat16):
+ assert torch.equal(F.scaled_dot_product_attention(q, k, v), TORCH_SDPA(q, k, v))
+ assert torch.equal(F.scaled_dot_product_attention(q, k, v), TORCH_SDPA(q, k, v))
+
+ def test_extra_memory_per_call_stays_within_the_measured_bound(self, kernel):
+ """Including the first call, whose check against PyTorch must not cost a full-size reference."""
+ q, k, v = (torch.randn(1, 4608, 24, 128, device="cuda", dtype=torch.bfloat16).transpose(1, 2) for _ in range(3))
+ kernel(q, k, v, tensor_layout="HND", is_causal=False, sm_scale=128**-0.5) # load the kernels unmeasured
+ # Quantized copies of Q, K and V at ~1 byte per element; measured 1.7x that on an RTX 4090.
+ bound = 2.5 * (q.numel() + k.numel() + v.numel())
+ with sage_attention_scope():
+ for call in ("first, checked", "later"):
+ torch.cuda.synchronize()
+ torch.cuda.reset_peak_memory_stats()
+ before = torch.cuda.memory_allocated()
+ out = F.scaled_dot_product_attention(q, k, v)
+ torch.cuda.synchronize()
+ extra = torch.cuda.max_memory_allocated() - before - out.numel() * out.element_size()
+ assert extra <= bound, call
+ del out
diff --git a/tests/test_config.py b/tests/test_config.py
index 4bcf5a9a444..3a1057ac6d9 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -117,6 +117,16 @@ def test_db_synchronous_defaults_to_full_and_loads_from_yaml(tmp_path: Path, pat
assert load_and_migrate_config(temp_config_file).db_synchronous == "normal"
+def test_attention_backend_defaults_to_auto_and_loads_from_yaml(tmp_path: Path, patch_rootdir: None) -> None:
+ # The default must stay `auto`: SageAttention changes images at the same seed, so it is opt-in only.
+ assert InvokeAIAppConfig().attention_backend == "auto"
+
+ temp_config_file = tmp_path / "temp_invokeai.yaml"
+ temp_config_file.write_text('schema_version: "4.0.3"\nattention_backend: sage\n')
+
+ assert load_and_migrate_config(temp_config_file).attention_backend == "sage"
+
+
def test_load_config_from_root_reads_that_roots_file_and_writes_nothing(tmp_path: Path, patch_rootdir: None) -> None:
root = tmp_path / "install"
root.mkdir()