[Feature Request] AMD Support

[0xDELUXA]

@woct0rdho @jammm any plans to add AMD support here, similar to Jam’s upstream PR thu-ml#108? I assume that won’t be merged there. I’m especially interested in RDNA4 support, which isn’t covered yet.

[woct0rdho]

I need to think about how to do it. I haven't tried it. The first purpose of this repo is to provide prebuilt wheels. It looks doable for me to merge Jam's code and build a wheel for some ROCm version with bundled kernels for gfx110x/115x/120x. The second purpose of this repo is to support Python stable ABI, libtorch stable ABI, and torch.compile . It also looks doable to me. The questions are:

1. What ROCm version should we build with? I don't think ROCm 7.2 is stable enough to support gfx1151 with ideal performance. Currently I'm using gfx1151 with some random nightly version.
2. Did you try that it works on RDNA4? Jam only says it supports RDNA3/3.5 .
3. Is it really the proper way to do sparse attention on RDNA3/3.5 ? We know that RDNA3/3.5 don't have faster int8 wmma than fp16, so I don't think int8 quant (the Sage part) is the proper way.

• I heard some reports that SageAttention (with int8, without sparse) is faster than torch SDPA on RDNA3/3.5, but I guess it would be actually because of a better pipeline rather than int8, and it would not be faster than a fp16 attention with well-written pipeline, such as the one in AITER.
• If SageAttention is really faster than AITER, that would be interesting. That should be because int8 saves LDS/VGPR consumption and bandwidth, and fp8 should work as well as int8 in this way with simpler quantization than int8. Last time I checked both the SageAttention in SageAttention AMD version thu-ml/SageAttention#332 and the SageAttention 1 Triton kernel are slower than AITER on my machine (gfx1151).

Also, in the long term I'm trying to move my wheel repos to https://github.com/Comfy-Org/wheels

[0xDELUXA]

1. What ROCm version should we build with? I don't think ROCm 7.2 is stable enough to support gfx1151 with ideal performance. Currently I'm using gfx1151 with some random nightly version.

In my opinion, TheRock should be used.

2. Did you try that it works on RDNA4? Jam only says it supports RDNA3/3.5.

It can achieve 0.99+ cosine similarity with my modified file from the last comment, but the performance is still lacking compared to SDPA Flash via AOTriton. I think this is because Sparge currently has no FP8 support.

3. Is it really the proper way to do sparse attention on RDNA3/3.5 ? We know that RDNA3/3.5 don't have faster int8 wmma than fp16, so I don't think int8 quant (the Sage part) is the proper way.

I'm not sure about this. Would be great to hear @jammm's opinion.

• I heard some reports that SageAttention (with int8, without sparse) is faster than torch SDPA on RDNA3/3.5, but I guess it would be actually because of a better pipeline rather than int8, and it would not be faster than a fp16 attention with well-written pipeline, such as the one in AITER.

SageAttention V1 on RDNA4 is better than SDPA for things like video generation.

• If SageAttention is really faster than AITER, that would be interesting. That should be because int8 saves LDS/VGPR consumption and bandwidth, and fp8 should work as well as int8 in this way with simpler quantization than int8. Last time I checked both the SageAttention in SageAttention AMD version thu-ml/SageAttention#332 and the SageAttention 1 Triton kernel are slower than AITER on my machine (gfx1151).

I don't really understand what you mean by SageAttention 1 Triton vs AITER. Did you mean FlashAttention-2's AITER Triton backend? I'm referring specifically to Windows here.

Great to see you're open to adding AMD support. I think the more options we have (e.g., attention types), the better. We already have SDPA, SageV1, and Flash-2 - the latter two are only possible thanks to your work on porting and maintaining Triton on Windows. It’d be great to have Sparge as well (based on SageAttention V2 via rocWMMA), if there’s already a way to support it.

[0xDELUXA]

I think I'll close this, as things are progressing well in thu-ml/SageAttention#368.

There's also woct0rdho/SageAttention#92 now.