From 6c8718945af881bdaeb2b092567df7e094370eee Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 3 Oct 2026 21:55:10 +0200 Subject: [PATCH] feat(z-image): run Z-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`, Z-Image denoising enters the SageAttention scope; Z-Image Control masks its attention and keeps SDPA. RTX 4090, Z-Image FP8: 10.4 s → 10.0 s at 1024², 33.7 s → 25.9 s at 2048², peak VRAM unchanged; a blind review judged no image worse. --- .../content/docs/configuration/Optimization/sage-attention.mdx | 1 + invokeai/app/invocations/z_image/z_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 4a0d5190e0c..915182f6cde 100644 --- a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx +++ b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx @@ -28,6 +28,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. | 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/z_image/z_image_denoise.py b/invokeai/app/invocations/z_image/z_image_denoise.py index 5894899b873..e9d7659d1b8 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 @@ -699,6 +700,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