From 4abdcd48dde9e59a371739a3646e786e6f1879fa Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 3 Oct 2026 21:55:11 +0200 Subject: [PATCH 1/2] feat(qwen-image): run Qwen-Image attention on SageAttention when enabled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With `attention_backend: sage`, Qwen-Image denoising enters the SageAttention scope. RTX 4090, Qwen-Image 2512 GGUF Q4_K_M: 78.8 s → 73.5 s at 1024², 208.9 s → 150.3 s at 2048², peak VRAM unchanged; a blind review judged no image worse. --- .../content/docs/configuration/Optimization/sage-attention.mdx | 1 + invokeai/app/invocations/qwen_image/qwen_image_denoise.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx index 915182f6cde..563925e6255 100644 --- a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx +++ b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx @@ -29,6 +29,7 @@ compared with PyTorch's default attention there: | 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. | 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. diff --git a/invokeai/app/invocations/qwen_image/qwen_image_denoise.py b/invokeai/app/invocations/qwen_image/qwen_image_denoise.py index 3cd1d897171..f161d20c6ce 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( @@ -484,6 +485,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) From 482e7f348f0fc5e1a0b0eeb8ce9ccf0ec3d0646a Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 3 Oct 2026 21:55:12 +0200 Subject: [PATCH 2/2] feat(krea2): run Krea-2 attention on SageAttention when enabled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With `attention_backend: sage`, Krea-2 denoising enters the SageAttention scope unless INVOKE_KREA2_SDPA_BACKEND pins a PyTorch kernel. RTX 4090, Krea-2 Turbo GGUF Q4_K_M: 11.5 s → 11.0 s at 1024², 42.7 s → 36.2 s at 2048², peak VRAM unchanged; a blind review judged no image worse. --- .../docs/configuration/Optimization/sage-attention.mdx | 1 + invokeai/app/invocations/krea2/krea2_denoise.py | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx index 563925e6255..dbbe1133297 100644 --- a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx +++ b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx @@ -30,6 +30,7 @@ compared with PyTorch's default attention there: | 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. | 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. diff --git a/invokeai/app/invocations/krea2/krea2_denoise.py b/invokeai/app/invocations/krea2/krea2_denoise.py index d9aca4e851b..b17c1066d14 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)):