Sage Attention 2 not work with Pytorch Compile

[pivtienduc]

Hi @jt-zhang

It seem that with the pytorch compile have not worked with the sage attention

I try to use just Flux basic ksampler workflow but get an full of noise image like this, the inference speed still increase but image is broken

This is my comfyUI:

When I turn off Sage Attention either pytorch cross attention or xformers, everything work fine:

I already test many time with many ComfyUI version and pytorch version, reinstall ComfyUI many times the problem still persist

Anything involve with Sage Attention and Torch Compile will create a noise image

However someone who own RTX4000s series will not encouter this problem. My gpu is 3090, is it a hardware related ?

Please support us, thank you so much

[woct0rdho]

What if you use SageAttention but don't use torch.compile?

[pivtienduc]

What if you use SageAttention but don't use torch.compile?

Then everything work fine,

And I found out that first time when I get a noise image. Disable sage attention then run with pytorch attention and let ComfyUI run pytorch compile then turn sage attention on, no more noise image but also there aren't any speed up with Torch Compile anymore, it run like no compile at all.

[StrongerXi]

This is strange, I've seen a few reports of sage attention not working with torch compile, I'll try look into this in the next week or so. You also mentioned Sage Attention 2 in your title, have you tried older version of sage attention?

[pivtienduc]

This is strange, I've seen a few reports of sage attention not working with torch compile, I'll try look into this in the next week or so. You also mentioned Sage Attention 2 in your title, have you tried older version of sage attention?

I haven't try sage attention 1 yet, ok let me check later then give you the result

[pivtienduc]

This is strange, I've seen a few reports of sage attention not working with torch compile, I'll try look into this in the next week or so. You also mentioned Sage Attention 2 in your title, have you tried older version of sage attention?

Well result is still the same. This is done on freshly install ComfyUI I just installed few hours ago

Please check

[BaseMe2]

Edit: Working on my end now.
I did downgrade sageattention to "sageattention-2.1.1+cu128torch2.7.0-cp311-cp311-win_amd64",
removed a kjnode, added --use-sage-attention flag.
Edit: Not fixed. I keep shift away completely now.In combination with using no "shift" node with a higher value than 1.15, as higher is for video generation only.

TorchCompile still doesn't like to run with kjnodes "TorchCompileModelFluxAdvanced", maybe I find a different way.

env_list.txt <- I'll leave it just in case.

[BaseMe2]

related: woct0rdho/triton-windows#120

[pikerbright]

@pivtienduc @StrongerXi I found that query, key, value input stride cause that sage attention can't work with torch compile.
following code can work.

def sageattention_forward(
    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,
) -> torch.Tensor:
    query = query.permute(0, 2, 1, 3)
    key = key.permute(0, 2, 1, 3)
    value = value.permute(0, 2, 1, 3)
    out = sageattn(
        query,
        key,
        value,
        attn_mask=attn_mask,
        dropout_p=dropout_p,
        is_causal=is_causal,
        scale=scale,
        tensor_layout='NHD'
    )
    return out.permute(0, 2, 1, 3)

[pivtienduc]

@pivtienduc @StrongerXi I found that query, key, value input stride cause that sage attention can't work with torch compile. following code can work.

def sageattention_forward(
    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,
) -> torch.Tensor:
    query = query.permute(0, 2, 1, 3)
    key = key.permute(0, 2, 1, 3)
    value = value.permute(0, 2, 1, 3)
    out = sageattn(
        query,
        key,
        value,
        attn_mask=attn_mask,
        dropout_p=dropout_p,
        is_causal=is_causal,
        scale=scale,
        tensor_layout='NHD'
    )
    return out.permute(0, 2, 1, 3)

Hi @StrongerXi

Can you take a look about it ?

Thank you so much

[woct0rdho]

Hi all, here is one more thing to consider: Did you launch ComfyUI with --fast?

@pivtienduc Looking from your log, there is Enabled fp16 accumulation.

I've found cases where --fast fp16_accumulation (or simply --fast which enables all fast options) causes the output to be noise, while --fast cublas_ops seems to be fine. (--fast fp8_matrix_mult doesn't work on RTX 30xx so I haven't tested it yet)

It's the same as OP described: When running the ComfyUI official Flux workflow (with GGUF model, but the problem is unlikely in GGUF) using SageAttention (launch with --use-sage-attention) and a compile node, the output is noise. When using only SageAttention, or only the compile node (with xformers attention), the output is fine.

I did not see significant speedup with --fast fp16_accumulation, so it's better to disable it in exchange for SageAttention + compile. BTW, --fast cublas_ops has some speedup for SDXL so I enable it.

I'm using torch 2.7.1+cu128, sageattention 2.1.1, triton-windows 3.3.1, cublas_ops 0.0.5, xformers 0.0.30, ComfyUI bd951a7

Looking from the reports here, this happens on RTX 30xx, 40xx, and 50xx.

[pivtienduc]

Hi all, here is one more thing to consider: Did you launch ComfyUI with --fast?

@pivtienduc Looking from your log, there is Enabled fp16 accumulation.

I've found cases where --fast fp16_accumulation (or simply --fast which enables all fast options) causes the output to be noise, while --fast cublas_ops seems to be fine. (--fast fp8_matrix_mult doesn't work on RTX 30xx so I haven't tested it yet)

It's the same as OP described: When running the ComfyUI official Flux workflow (with GGUF model, but the problem is unlikely in GGUF) using SageAttention (launch with --use-sage-attention) and a compile node, the output is noise. When using only SageAttention, or only the compile node (with xformers attention), the output is fine.

I did not see significant speedup with --fast fp16_accumulation, so it's better to disable it in exchange for SageAttention + compile. BTW, --fast cublas_ops has some speedup for SDXL so I enable it.

I'm using torch 2.7.1+cu128, sageattention 2.1.1, triton-windows 3.3.1, cublas_ops 0.0.5, xformers 0.0.30, ComfyUI bd951a7, RTX 3080.

Thank you so much. I will try --fast cublas_ops

[woct0rdho]

@jt-zhang Maybe in SageAttention we can add a warning if torch.backends.cuda.matmul.allow_fp16_accumulation is enabled