diff --git a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx index 910590ffe04..e309c7f3d03 100644 --- a/docs/src/content/docs/configuration/Optimization/sage-attention.mdx +++ b/docs/src/content/docs/configuration/Optimization/sage-attention.mdx @@ -27,6 +27,7 @@ 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 %) | | 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/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