RDNA4 owners on ComfyUI: I made SageAttention actually fast on the 9070 XT, here's the build plus every caveat
TL;DR: A drop-in sageattention build for RX 9070 / 9070 XT on Windows. The common case (fp16, head_dim 128) runs on an fp8 attention kernel I wrote by hand in HIP. Against the existing gfx12 port (SageAttention PR #368), each attention call is 1.07–1.21× faster non-causal and 1.72–1.96× faster causal. In ComfyUI with Krea2 that works out to a 5–14% faster sampling step than PyTorch SDPA. Install the wheel, keep using --use-sage-attention, done. Caveats below, and there are a few, so please read them.
Repo + wheel: https://github.com/IxMxAMAR/SageAttention-RDNA4
Why I did this
I run ComfyUI on a 9070 XT under Windows. SageAttention on RDNA4 works thanks to the gfx12 port in PR #368, so huge credit to DELUXA for that. But I wanted to see how far this card can actually go. So I started with a Triton kernel, lost to #368 on most shapes, spent a long time figuring out why, and ended up writing the kernel by hand.
What finally made the difference:
- 64 keys per loop iteration instead of 16.
- Computing Q·Kᵀ transposed. The result then lands in exactly the register layout the next matrix multiply needs, so the attention weights never leave registers.
- One scale per token for Q, K and V, instead of one per 64-token block.
The whole story, including everything that didn't work (most things didn't, lol), is in the repo's docs/JOURNEY.md and docs/FINDINGS.md.
Who this is for, right now
- ComfyUI users on an RX 9070 or RX 9070 XT, on Windows. This is what I've tested and what I use daily.
- Anyone else who calls
sageattnfrom PyTorch on the same card should be able to use it too; it's the same package API. But I haven't tested any other apps or workflows yet. If you try it outside ComfyUI, tell me how it goes.
Speed (image 1)
Same tensors, same process, interleaved runs, full call including quantization:
- Non-causal (what image/video diffusion uses): 1.07–1.21× faster than #368.
- Causal: 1.72–1.96× faster.
- These are kernel-level numbers, and kernel-level isn't what you feel. Image 2 is what you feel.
In an actual ComfyUI render (image 2)
Krea2, 8 steps, fp16, called exactly the way stock ComfyUI calls it:
| Backend | 1 MP | 2 MP |
|---|---|---|
| PyTorch SDPA | 1.03 s/step | 2.39 s/step |
| SageAttention 1.x | 1.00 s/step | 2.17 s/step |
| This build | 0.97 s/step | 2.06 s/step |
Attention itself is about 2× faster than SDPA per call. But most of a step is the model's matmuls, so the end-to-end gain is 5–14% over SDPA and 2–5% over SageAttention 1.x. The gain grows with resolution.
Accuracy, honestly (image 3)
Everything is fp8 with per-token scales. I measured against exact attention on real Q/K/V captured from inside Krea2 and Flux2-Klein:
- On "calm" layers, #368 is more precise. It uses int8 for Q·K, which has finer steps than fp8.
- On layers with a few extreme keys, #368's error blows up, by up to ~200×. Krea2's first block is one such layer. One outlier wrecks the scale for its whole 64-token block. Per-token scaling never has that bad case.
- smooth_k is on by default and should stay on. It made my kernel more accurate on 10 out of 10 real captures.
In my Krea2 tests the images were closer to full-precision attention than with SageAttention 1.x. I can't tell them apart by eye. Like any quantized attention, though, your images won't match SDPA pixel for pixel: a turbo 8-step sampler amplifies tiny differences into different details.
Caveats (please read)
- RX 9070 / 9070 XT (gfx1201) only. On any other GPU it falls back to #368's kernels: no harm, no gain.
- Windows only, as far as testing goes. Linux is untested.
- The prebuilt wheel needs exactly Python 3.12 + PyTorch
2.13.0+rocm10.0.0**.** The compiled parts are tied to that PyTorch version. Anything else means building from source, which takes about 50 minutes. - fp16 models only for now. bf16 models silently fall back to #368's kernel. Still works, just no speed-up from my kernel.
- head_dim 128 only. That covers Flux-style models such as Krea2, Flux and Klein. Other head sizes fall back.
- Tested end to end on Krea2 only. The kernel accepts the input layout video models like Wan use, and it's checked bit for bit, but I haven't benchmarked a video workflow yet.
- Very small images (well under 1 MP) can come out a hair slower than SageAttention 1.x. The win is at normal and high resolutions.
- No attention masks. Masked calls go to the fallback.
- LLMs: only prompt processing (prefill) through PyTorch would benefit. Token-by-token generation isn't covered, and llama.cpp / LM Studio don't use this at all.
- One bug already found and fixed. A pre-release build could turn a render black when a single value got huge enough to overflow fp8. It's fixed in this release, and every conversion in the kernel is now stress-tested with extreme values.
Install
- Close ComfyUI.
- Download the wheel from the GitHub release.
path\to\ComfyUI\venv\Scripts\python -m pip install --no-deps sageattention-2.2.0+amd.gfx12.1-cp312-cp312-win_amd64.whl(portable ComfyUI: usepython_embeded\python.exeinstead)- Start ComfyUI with
--use-sage-attentionlike always.
--no-deps matters: it stops pip from touching your PyTorch. To switch my kernel off without uninstalling, set SAGEATTN_SK1_BACKEND=0.
What's next
- bf16 support, so bf16 models get the speed-up too. This is the big one for coverage.
- int8 Q·K with per-token scales, aiming for #368's precision on calm layers and no blow-ups on outlier layers. Already in the works.
- head_dim 64, for SDXL-era models. Also in the works.
- An end-to-end test on a video workflow (Wan), since the layout support is already in.
- An experiment to overlap the two halves of the kernel, which might be worth another few percent. A quick test decides whether it's worth building.
If you've got a 9070 / 9070 XT, I'd love numbers from your setup: your model, resolution and s/it before/after. Bug reports are very welcome too. Credit to the SageAttention team (thu-ml) and to DELUXA for #368; this builds directly on their work. Apache-2.0.