sageattn_ultravico error in ComfyUI

[0xDELUXA]

Describe the bug

Using Kijai's sageattn_ultravico as the attention mode for ksampler in ComfyUI, the following error occurs:
https://gist.github.com/0xDELUXA/a80856f7071894d0cf777d7e93ba5935

It looks similar to #184, but f5ea393 did not fix it.
Only ultravico fails; sageattn, radial_sage_attention, and sageattn_compiled all work correctly (after applying the patches below).

Could this actually be an issue with sageattn_ultravico itself, rather than triton-windows?

Environment details

GPU: AMD Radeon RX 9060 XT
Python: 3.12
PyTorch: 2.11.0a0+rocm7.11.0a20260121
Triton: 3.6.0a0.post25

[woct0rdho]

It's a different bug from #184

[0xDELUXA]

I also thought that the issue might be in ComfyUI-WanVideoWrapper\ultravico\sageattn\attn_qk_int8_per_block.py. I assume it works on AMD (on Linux), so this could be a Windows-specific bug. AFAIK, only ultravico uses its own attn_qk_int8_per_block.py, that's in ComfyUI\custom_nodes\ComfyUI-WanVideoWrapper\ultravico\sageattn. The other sage implementations use ComfyUI\venv\Lib\site-packages\sageattention\attn_qk_int8_per_block.py

[rwfsmith]

the base SageAttention 1.0.6 package with the comfyui native nodes gives the same error.

[0xDELUXA]

You're right. I completely forgot that sageattention 1.0.6 doesn't work without these community-made patches:

curl -s -o venv\Lib\site-packages\sageattention\attn_qk_int8_per_block.py https://raw.githubusercontent.com/patientx/ComfyUI-Zluda/refs/heads/master/comfy/customzluda/sa/attn_qk_int8_per_block.py 
curl -s -o venv\Lib\site-packages\sageattention\attn_qk_int8_per_block_causal.py https://raw.githubusercontent.com/patientx/ComfyUI-Zluda/refs/heads/master/comfy/customzluda/sa/attn_qk_int8_per_block_causal.py 
curl -s -o venv\Lib\site-packages\sageattention\quant_per_block.py https://raw.githubusercontent.com/patientx/ComfyUI-Zluda/refs/heads/master/comfy/customzluda/sa/quant_per_block.py venv\\Lib\site-packages\sageattention\quant_per_block.py

[0xDELUXA]

It’d be great to have SageAttention working OOTB on AMD/Windows. And I know who could make it happen: @jammm

[rwfsmith]

Those files fixed my error, but it doesn't seem any faster than Flash Attention 2 for me, if not very slightly slower, and with a noticeable loss in quality.

[woct0rdho]

@rwfsmith Are you using Strix Halo? RDNA3 GPUs do not have faster int8 than fp16, so SageAttention will not be faster than FlashAttention on them. On RDNA4 GPUs there should be a way to make SageAttention faster than FlashAttention.

[rwfsmith]

@rwfsmith Are you using Strix Halo? RDNA3 GPUs do not have faster int8 than fp16, so SageAttention will not be faster than FlashAttention on them. On RDNA4 GPUs there should be a way to make SageAttention faster than FlashAttention.

Yeah, it's Strix Halo. I've got a 5090 for my diffusion needs, but I still like to fiddle with it on the Halo to see if I can get stuff working :P. I mostly use it for LLMs. Training is better now with FA2 working.

[0xDELUXA]

@rwfsmith Are you using Strix Halo? RDNA3 GPUs do not have faster int8 than fp16, so SageAttention will not be faster than FlashAttention on them. On RDNA4 GPUs there should be a way to make SageAttention faster than FlashAttention.

Theoretically, yes. Right now, using those patches on gfx1200, Sage V1 gives more or less the same performance and memory usage as FA-2. I'm currently waiting for Jam's SpargeAttention PR to be compatible with RDNA4 - it's based on eliotwang's Sage V2 implementation, so it should be even faster for us RDNA4 users.

[0xDELUXA]

Those files fixed my error, but it doesn't seem any faster than Flash Attention 2 for me, if not very slightly slower, and with a noticeable loss in quality.

Those patches were made for ZLUDA before ROCm was available on Windows and originally targeted older architectures.

[woct0rdho]

Is it a Windows-specific bug? If on Linux we also need the patches from ComfyUI-Zluda to make the SageAttention Triton kernels run, then it should be a bug in the AMD backend rather than the Windows support code in Triton (I haven't tried)

[rwfsmith]

I can boot up Ubuntu to test it out later.

[0xDELUXA]

Is it a Windows-specific bug? If on Linux we also need the patches from ComfyUI-Zluda to make the SageAttention Triton kernels run, then it should be a bug in the AMD backend rather than the Windows support code in Triton (I haven't tried)

I think so too. If it doesn't work on Linux either, then this issue doesn't really belong to the triton-windows repo. In that case, I'll close it.