YAML input for Sage Attention node? #12

EnragedAntelope · 2024-12-01T02:48:57Z

EnragedAntelope
Dec 1, 2024

Apologies if I missed it but I read the readme and didn't see info on what the YAML input is expecting specifically for the Sage Attention new nodes? Please share guidance on how best to use this... I tried SageAttn2 with CogVideoX via Kijai's nodes there and am very impressed at the speedup so looking forward to using it more. Thanks!

blepping · 2024-12-01T06:48:08Z

blepping
Dec 1, 2024
Maintainer

No need to apologize. There isn't really specific documentation for using it as it is an advanced feature that allows you to directly pass (or override) arguments for the SageAttention function(s) which are provided by the SageAttention package. So basically, if something in the package changes, that would also change what you can/should pass.

This reply got pretty long so if you want a TL;DR: If you don't already know you need to change something there then it's probably not something you need to worry about/mess with. If you still want to know what you could use it to do, then read on!

You pretty much need to look at the source code for SageAttention to see what it parameters you can change: https://github.com/thu-ml/SageAttention/blob/main/sageattention/core.py

The yaml_parameters input can be connected to a multiline text widget or converted into one in the node itself. Here is an example of parameters you could possibly pass (to the v2.0.0 beta version):

Most keys you specify will just be passed to the function directly but there are a couple exceptions. sageattn_function can be used to select a function that the SageAttention package provides by name so you could use it to override the kernel SageAttention automatically selected for your architecture. Should you do this? Likely not, but you can if you want to. sm_scale is the softmax scale (don't ask me exactly what it does because I don't know!) which normally automatically gets set to the tensor_head_dimension ** -0.5 but you can mess with it to get a different effect.

Just for example, the only head dimension in SDXL that SageAttention currently supports is 64, 64 ** -0.5 == 0.125 so for SDXL values close to 0.125 will be roughly what would be auto-selected. Raising/lowering it a bit seems to produce effects a bit like CFG, a bit like ComfyUI's experimental attention multiplier nodes. YAML parameters currently has one other special key you can pass, and that is setting sm_scale_<head dim here>, for example sm_scale_64: 0.125 would use that scale specifically for attentions with head dimension 64 and you could set sm_scale_96: <other value> if you wanted to use something else for a different head dimension.

Note that I'm not saying it's better to change that or anything, just that you can for a different effect (I like having knobs to turn!). Passing incorrect parameters, forcing a kernel your architecture doesn't expect, etc are all things that are likely to cause SageAttention to have unintended behavior/perform worse/maybe even crash ComfyUI or Torch. In other words, if you mess with that then you shouldn't bug the SageAttention authors if you run into problems.

Hope this helps answer your question!

0 replies

blepping · 2024-12-01T06:57:14Z

blepping
Dec 1, 2024
Maintainer

By the way, I recommend using SageAttentionSampler rather than GlobalSageAttention when possible. You're less likely to run into problems (since my implementation patches/replaces ComyfUI's default optimized attention implementation). Also, the global node isn't really a model patch, so for example if you were to bypass it that wouldn't actually disable SageAttention: you have to actually make sure the node runs with your changed settings (including the enabled/disabled toggle) for them to have effect.

The sampler specific version has less sharp edges.

0 replies

EnragedAntelope · 2024-12-03T01:29:55Z

EnragedAntelope
Dec 3, 2024
Author

Thank you for the helpful explanation.
I just had a chance to retry this. I used the sampler this time and it works fine and did not force me to add YAML parameters, so I skipped that entirely. However I saw no noticeable speed difference with flux at 896x1152, is that expected? It sounds like it would only kick in with a hiresfix sort of node I suppose?

1 reply

blepping Dec 3, 2024
Maintainer

I used the sampler this time and it works fine and did not force me to add YAML parameters, so I skipped that entirely.

Oh, sorry and thanks for point that out. yaml_parameters was supposed to be optional for both nodes. Should be fixed now.

However I saw no noticeable speed difference with flux at 896x1152, is that expected?

As far as I know, SDXL is the only image model with compatible attentions (some video models might be, I haven't tested them). It didn't seem to do anything for Flux when I tested, it definitely doesn't work with SD15.

You can try that sageattn_verbose: true parameter show in the screenshot above if you want to double check it. It should spam your console with messages like:

>>> SAGE(False): reshape=..... blah blah blah, dim_head=<some number>

SAGE(False) means the attention was not supported so it just got delegated to non-SageAttention. This is probably what you'll see, if you see SAGE(True) that means it's actually using SageAttention. Looking at dim_head might also be useful: SageAttention only supports 64, 96 and 128 at the moment. They say they're working on supporting other dimensions, so once that is added I will update the node to be able to use them. Unfortunately, it's not really something I can do myself: have to wait for them to support it, if they ever do.

blepping · 2024-12-28T09:35:12Z

blepping
Dec 28, 2024
Maintainer

@EnragedAntelope I just merged some changes for the bleh SageAttention support that should make it work for Flux now (in my limited testing I didn't notice a massive performance increase though). I recommend using the recently released SageAttention 2.0.1 if possible. Will also work for some other models like SD 1.5 which it didn't before.

0 replies

EnragedAntelope · 2024-12-28T11:47:17Z

EnragedAntelope
Dec 28, 2024
Author

Thank you! FWIW Kijai nodes also do this so now two sage attn options for flux/other.

0 replies

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

YAML input for Sage Attention node? #12

{{title}}

Replies: 5 comments 1 reply

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

{{title}}

Select a reply

YAML input for Sage Attention node? #12

EnragedAntelope Dec 1, 2024

Replies: 5 comments · 1 reply

blepping Dec 1, 2024 Maintainer

blepping Dec 1, 2024 Maintainer

EnragedAntelope Dec 3, 2024 Author

blepping Dec 3, 2024 Maintainer

blepping Dec 28, 2024 Maintainer

EnragedAntelope Dec 28, 2024 Author

EnragedAntelope
Dec 1, 2024

Replies: 5 comments 1 reply

blepping
Dec 1, 2024
Maintainer

blepping
Dec 1, 2024
Maintainer

EnragedAntelope
Dec 3, 2024
Author

blepping Dec 3, 2024
Maintainer

blepping
Dec 28, 2024
Maintainer

EnragedAntelope
Dec 28, 2024
Author