How to use sageattn in flux?

[digdigdigging]

I tried to modify the original code of diffusers, in FluxAttnProssor, I set '_attention_backend' to '_sage_qk_int8_pv_fp16_cuda', which goes to attention_dispatch.py(code as below). It worked, but the output image turn into noise picture. How exactly to use sageattn in flux?

@_AttentionBackendRegistry.register(
AttentionBackendName._SAGE_QK_INT8_PV_FP16_CUDA,
constraints=[_check_device_cuda_atleast_smXY(8, 0), _check_shape],
)
def _sage_qk_int8_pv_fp16_cuda_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
is_causal: bool = False,
scale: Optional[float] = None,
qk_quant_gran: _SAGE_ATTENTION_QK_QUANT_GRAN = "per_thread",
pv_accum_dtype: _SAGE_ATTENTION_PV_ACCUM_DTYPE = "fp32",
smooth_k: bool = True,
smooth_v: bool = False,
return_lse: bool = False,
) -> torch.Tensor:
return sageattn_qk_int8_pv_fp16_cuda(
q=query,
k=key,
v=value,
tensor_layout="NHD",
is_causal=is_causal,
qk_quant_gran=qk_quant_gran,
sm_scale=scale,
pv_accum_dtype=pv_accum_dtype,
smooth_k=smooth_k,
smooth_v=smooth_v,
return_lse=return_lse,
)

[digdigdigging]

Or I directly change the native attn function in attention_dispatch.py , the output image still turn into noise picture, why?

@_AttentionBackendRegistry.register(
AttentionBackendName.NATIVE,
constraints=[_check_device, _check_shape],
)
def _native_attention(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attn_mask: Optional[torch.Tensor] = None,
dropout_p: float = 0.0,
is_causal: bool = False,
scale: Optional[float] = None,
enable_gqa: bool = False,
) -> torch.Tensor:
query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
# out = torch.nn.functional.scaled_dot_product_attention(
# query=query,
# key=key,
# value=value,
# attn_mask=attn_mask,
# dropout_p=dropout_p,
# is_causal=is_causal,
# scale=scale,
# enable_gqa=enable_gqa,
# )
out = sageattn(...)
out = out.permute(0, 2, 1, 3)
return out

[anm-ol]

https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends#setattentionbackend
you just need to put
pipe.transformer.set_attention_backend("sage_hub") before calling the pipeline