gfx1201 (RDNA4) SageAttention attention backend: aiter package or Hub kernel? #15014

Is your feature request related to a problem? Please describe.

On AMD Radeon RDNA4 GPUs (gfx1201: Radeon AI PRO R9700 / R9600D, RX 9070 series), diffusers has no quantized attention backend:

• SAGE and SAGE_HUB are CUDA builds.
• The ROCm Hub FlashAttention kernel (kernels-community/aiter-flash-attn) ships tuning configs for gfx942 / gfx950 / gfx1250 only.

So these GPUs use native SDPA.

aiter has a SageAttention kernel for gfx1201: aiter.gfx1201_sage_attention(q, k, v, softmax_scale), ROCm/aiter#6390 (under review).

• INT8 QK^T and FP8 PV with FP32 accumulation; Q/K are quantized per 32-row block and V per channel inside the kernel.
• Inputs: BF16 [B, S, H, 128], Q/K/V of equal shape, no mask, non-causal.

MiniMax-H3 in diffusers matches that contract exactly. MiniMaxH3AttnProcessor passes no mask, the attention is self-attention with head dim 128, and the input is BF16. The same applies to the token refiner.

For reference, the same kernel as an opt-in vLLM-Omni backend (vllm-project/vllm-omni#8792, draft) on MiniMax-H3 T2VA: 2 x R9700 TP2, 1344x768, 4 s, 8 steps.

Describe the solution you'd like.

An opt-in diffusers attention backend for gfx1201 SageAttention. It would never be selected by default, and it would reject unsupported inputs (masks, causal, head dim other than 128, unequal Q/K/V shapes) with a clear error.

It would support Ulysses context parallelism (inference only, no LSE needed), so MiniMax-H3 can run with its existing _cp_plan on 2+ GPUs.

Before writing it, I'd like to check which delivery you prefer, given that #14436 moved ROCm FlashAttention from the aiter pip package to a Hub kernel:

1. AITER_SAGE, using the aiter package. The backend imports aiter and calls the op. It is a small diffusers change, but it adds back an optional aiter dependency and needs an aiter release that contains [HIP] [JIT] [Build] gfx1201: SageAttention forward (INT8 QK, FP8 PV) ROCm/aiter#6390.
2. AITER_SAGE_HUB, a Hub kernel. We publish the kernel as a kernels package. Like kernels-community/aiter-flash-attn, it would be a universal torch-rocm build: precompiled gfx1201 code objects plus a small Python loader registered as a torch custom op. diffusers would only add a _HubKernelConfig entry and the backend function. This does not depend on an aiter release. The open question is which Hub repo it should live in.
3. A ROCm variant of kernels-community/sage-attention, so the existing SAGE_HUB works unchanged. Its sageattn would route gfx1201 inputs it supports to this kernel and the rest to SDPA.

Describe alternatives you've considered.

• Keep SDPA on gfx1201. That leaves about 1.8x on MiniMax-H3 attention unused.
• A custom attention processor in user code. It works, but every model needs its own processor, and it does not compose with set_attention_backend or context parallelism.

Additional context.

I can implement whichever option you prefer, with tests and a MiniMax-H3 end-to-end comparison against SDPA on 2 x R9700 using Ulysses CP. A sparse variant (Sol-Attn, ROCm/aiter#6292) also exists. Following the note in attention_dispatch.py about per-model-tuned sparse methods, I am not proposing it here.

cc @sayakpaul

AI assistance: this proposal was drafted with the help of an AI coding assistant (GitHub Copilot). The numbers above were measured on the listed hardware.