feat(krea2): run Krea-2 attention on SageAttention when enabled

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.

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

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)