diff --git a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx index 66a2aeca8b1..4a632c09856 100644 --- a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx +++ b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx @@ -29,6 +29,8 @@ 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. | +| 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 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)