aiter-sage-attention: SageAttention for AMD RDNA4 (gfx1201) #1222

Hi! I'd like to propose hosting a SageAttention kernel for AMD Radeon RDNA4 GPUs (gfx1201: Radeon AI PRO R9700 / R9600D, RX 9070 series) here, and I'd be grateful for your thoughts.

The kernel

It computes QK^T in INT8 and PV in FP8, with FP32 accumulation. Q/K are quantized per 32-row block and V per channel inside the kernel. The source is in ROCm/aiter (HIP, MIT): ROCm/aiter#6390. That PR is still under review.

Why it could be useful here

As far as I can tell, diffusers doesn't have a quantized attention backend for these GPUs yet. SAGE/SAGE_HUB are CUDA builds, and kernels-community/aiter-flash-attn ships tuning configs for gfx942 / gfx950 / gfx1250. In huggingface/diffusers#15014, @sayakpaul suggested packaging the kernel with Kernels before adding a diffusers backend, and @Abdennacer-Badaoui kindly offered to help with the build.

What I have so far

A package laid out like aiter-flash-attn-ck (edition 5, rocm-archs = ["gfx1201"]):

• The two aiter kernel sources (prepare and attention) are vendored unchanged from aiter#6390 at 62019a49. aiter_tensor.h is trimmed to the fields the kernels use, so ck_tile isn't needed.
• It exposes sageattn(q, k, v, tensor_layout, is_causal, sm_scale, return_lse) for diffusers, plus gfx1201_sage_attention(query, key, value, softmax_scale, ...). Unsupported inputs (causal, LSE, non-BF16, head dim other than 128, unequal Q/K/V shapes, non-gfx1201) raise an error rather than falling back silently.
• tests/ loads the kernel via get_kernel.

I don't have Nix on my machine, so I built it locally with kernel-builder create-pyproject + setup.py build_kernel (PyTorch 2.11, ROCm 7.15). On an R9700:

• All 22 tests pass.
• Outputs are bit-identical to aiter.gfx1201_sage_attention built from the same commit.
• Compared with SDPA (BF16, 28 heads, head dim 128; cosine similarity 0.9992 to SDPA):

I also tried it end to end in diffusers, through a local aiter_sage_hub backend (loaded via LOCAL_KERNELS), with MiniMax-H3 text-to-video+audio (960x544, 124 frames, 9 steps, one R9700, group offload). Compared with the native backend:

• Each transformer call took 12.0 s instead of 16.5 s (1.37x), and the pipeline call took 130 s instead of 171 s.
• Video PSNR to the native output was 26.1 dB (worst frame 22.0 dB), and the two videos look the same to me.
• The VAE's attention (head dim 64) also follows the active backend, so the diffusers backend sends calls the kernel can't take to SDPA.

My question

Would you be open to hosting this as kernels-community/aiter-sage-attention? If so, could I be vouched so that I can open a PR? I'd open it as a draft and re-sync the vendored sources once aiter#6390 is merged. I'm also happy to adjust the layout or naming to whatever you prefer.

For transparency: I wrote the package with help from an AI coding assistant (GitHub Copilot). I reviewed the code and ran the tests and measurements above myself.

Thank you for your time!