Add `aiter_sage_hub` attention backend for AMD RDNA4 (gfx1201) by Shan2L · Pull Request #15060 · huggingface/diffusers

Shan2L · 2026-10-16T01:10:16Z

What does this PR do?

Adds an aiter_sage_hub attention backend for AMD RDNA4 (gfx1201) GPUs, as discussed in #15014. It runs SageAttention (INT8 QK^T, FP8 PV) from the kernels-community/aiter-sage-attention Hub kernel (huggingface/kernels-community#1229), which vendors the gfx1201 kernels from ROCm/aiter#6390.

Fixes #15014

pipe.transformer.set_attention_backend("aiter_sage_hub")

The kernel takes bf16 Q/K/V of equal shape with head dim 128, without a mask and non-causal. Other calls fall back to native SDPA instead of raising, because components such as the VAE (head dim 64) and cross-attention follow the active backend.
Ulysses context parallelism is supported, with the same fallback. Ring attention raises NotImplementedError, since it needs the LSE and the kernel does not return it.
Inference only, like the other Sage backends.

Results

MiniMax-H3 text-to-video+audio on one Radeon AI PRO R9700: 960x544, 124 frames, 9 steps (8 transformer calls), seed 42, torch 2.13, Hub kernel v1, transformer with block-level group offloading.

Backend	Transformer call (mean of calls 2-8)
`native`	16.76 s
`aiter_sage_hub`	12.10 s (1.39x)

Video PSNR to native is 23.3 dB (18.9 to 25.9 dB per frame) and audio SNR is 8.5 dB. For scale, the native backend alone differs by 23.7 dB video PSNR and 12.9 dB audio SNR between torch 2.11 and torch 2.13.

With Ulysses degree 2 on two GPUs, dispatch_attention_fn gives bit-identical output to the single-GPU run, both for self-attention (kernel) and for cross-attention with shorter K/V (fallback).

Benchmark setup

import torch
from diffusers import ModularPipeline
from diffusers.hooks import apply_group_offloading

pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-H3", workflow="t2va")
pipe.load_components(dtype=torch.bfloat16)
offload = dict(onload_device=torch.device("cuda"), offload_device=torch.device("cpu"), use_stream=True)
pipe.transformer.enable_group_offload(offload_type="block_level", num_blocks_per_group=1, **offload)
apply_group_offloading(pipe.text_encoder.model, offload_type="leaf_level", **offload)
pipe.vae.to("cuda")
pipe.audio_vae.to("cuda")
pipe.transformer.set_attention_backend("aiter_sage_hub")  # skipped for native

res = pipe(
    prompt="A red fox trotting through a snowy pine forest, snow crunching underfoot",
    height=544,
    width=960,
    num_frames=124,
    num_inference_steps=9,
    generator=torch.Generator().manual_seed(42),
    output=["videos", "audio", "sampling_rate"],
)

Transformer calls were timed with forward pre/post hooks and torch.cuda.synchronize().

Tests

$ python -m pytest tests/models/test_attention_dispatch.py -k AiterSageHub -v
tests/models/test_attention_dispatch.py::TestAiterSageHubFallback::test_supported_inputs_use_kernel PASSED
tests/models/test_attention_dispatch.py::TestAiterSageHubFallback::test_unsupported_inputs_fall_back_to_native[head_dim_64] PASSED
tests/models/test_attention_dispatch.py::TestAiterSageHubFallback::test_unsupported_inputs_fall_back_to_native[fp32] PASSED
tests/models/test_attention_dispatch.py::TestAiterSageHubFallback::test_unsupported_inputs_fall_back_to_native[cross_attention] PASSED
tests/models/test_attention_dispatch.py::TestAiterSageHubFallback::test_unsupported_inputs_fall_back_to_native[causal] PASSED
tests/models/test_attention_dispatch.py::TestAiterSageHubFallback::test_unsupported_inputs_fall_back_to_native[mask] PASSED
tests/models/test_attention_dispatch.py::TestAiterSageHubFallback::test_ring_attention_raises PASSED
7 passed, 14 deselected

$ make quality
All checks passed!
2138 files already formatted

$ python utils/check_copies.py && python utils/check_dummies.py

The tests use a stub in place of the Hub kernel, since CI has no gfx1201. The results above ran on gfx1201 with the published kernel.

Self-review notes

Final round of the self-review skill:

Blocking: none. Fixed during review: under context parallelism, calls the kernel cannot take (such as cross-attention) reached the kernel and raised. Both paths now share the same support check.
Non-blocking, left as is:
1. Unsupported calls fall back to SDPA instead of raising, unlike the other Hub backends, since the VAE and cross-attention follow the active backend.
2. The tests replace the Hub kernel with a stub. testing.md discourages call-level test doubles, but CI has no gfx1201. The fallback tests compare against the real native backend.
3. Not added to _ALL_BACKEND_PARAMS, same as aiter_fa2_hub: CI has no gfx1201, and the test models don't use head dim 128, so only the fallback would run.

Before submitting

Who can review?

@sayakpaul @Abdennacer-Badaoui

SageAttention (INT8 QK^T, FP8 PV) for gfx1201 through the kernels-community/aiter-sage-attention Hub kernel. Calls the kernel cannot take (head dim other than 128, non-bf16, masks, causal, unequal Q/K/V shapes) fall back to native SDPA, since components such as the VAE follow the active backend. Supports Ulysses context parallelism. Signed-off-by: Sylvan Liu <Sylvan.Liu@amd.com>

sayakpaul · 2026-10-16T08:09:25Z

+    constraints=[_check_device_cuda, _check_shape],
+    supports_context_parallel=True,
+)
+def _aiter_sage_attention_hub(


Have we tested the CP bit?

Yes. With Ulysses degree 2 on two gfx1201 GPUs, the output of dispatch_attention_fn is bit-identical to the single-GPU run, both for self-attention (kernel) and for cross-attention with shorter K/V (fallback). Ring attention raises NotImplementedError. TestContextParallelAttentionBackward checks gradients, so it doesn't apply to this inference-only backend.

sayakpaul · 2026-10-16T08:12:16Z

+
+
+@is_attention
+class TestAiterSageHubFallback:


Looks good but it needs to be decorated with something like is_aiter_sage_hub_supported?

These tests replace the Hub kernel with a stub, so they run on CPU without a gfx1201 GPU or kernels, and only check which calls go to the kernel and which fall back. Would you like a GPU test with the real kernel as well, skipped unless a gfx1201 GPU and kernels are available?

I think that would be nice? We could add that to

diffusers/tests/models/testing_utils/attention.py

Line 438 in b800f82

@pytest.mark.parametrize("backend", _ALL_BACKEND_PARAMS)

WDYT?

Done in a09f58b: aiter_sage_hub is now in _ALL_BACKEND_PARAMS, skipped unless a gfx1201 GPU and kernels are available. The Flux and QwenImage tiny models use head dims below 128, so they only exercise the native fallback. I also added TestAiterSageHubKernel, which runs the real kernel on head dim 128 and compares it against SDPA (cosine similarity 0.9993).

On a gfx1201 GPU (torch 2.13, kernels 0.16.1):

test_attention_dispatch.py -k AiterSage: 8 passed

Flux and QwenImage -k aiter_sage_hub: 8 passed, including test_compile

sayakpaul

Great stuff! Can we see some visual examples as well?

Shan2L · 2026-10-16T13:41:16Z

Great stuff! Can we see some visual examples as well?

Thanks! Here is MiniMax-H3 text-to-video+audio from the benchmark in the description (960x544, 124 frames, 9 steps, seed 42), native SDPA on the left and aiter_sage_hub on the right. The audio is from the aiter_sage_hub run. Both original videos, each with its own audio, are attached too.

t213_native.mp4

t213_aiter_sage_hub.mp4

HuggingFaceDocBuilderDev · 2026-10-17T10:25:50Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Add aiter_sage_hub to the attention backend tester parameters, skipped unless a gfx1201 GPU and `kernels` are available. The tiny test models use head dims below 128, so they only exercise the native fallback; TestAiterSageHubKernel checks the real kernel on head dim 128 against SDPA. Signed-off-by: Sylvan Liu <Sylvan.Liu@amd.com>

sayakpaul

Superb work!

github-actions Bot added documentation Improvements or additions to documentation models tests fixes-issue size/M PR with diff < 200 LOC labels Oct 16, 2026

sayakpaul reviewed Oct 16, 2026

View reviewed changes

sayakpaul approved these changes Oct 16, 2026

View reviewed changes

Merge branch 'main' into feat/aiter-sage-hub

c5207e3

sayakpaul added the performance Anything related to performance improvements, profiling and benchmarking label Oct 17, 2026

sayakpaul added this to the Release 0.42.0 milestone Oct 17, 2026

github-actions Bot added size/L PR with diff > 200 LOC and removed size/M PR with diff < 200 LOC labels Oct 19, 2026

sayakpaul approved these changes Oct 19, 2026

View reviewed changes

Shan2L mentioned this pull request Oct 20, 2026

[HIP] [JIT] [Build] gfx1201: SageAttention forward (INT8 QK, FP8 PV) ROCm/aiter#6390

Open

Navigation Menu

Repository navigation

Add `aiter_sage_hub` attention backend for AMD RDNA4 (gfx1201) - #15060

Add `aiter_sage_hub` attention backend for AMD RDNA4 (gfx1201)#15060
Shan2L wants to merge 3 commits into
huggingface:mainfrom
Shan2L:feat/aiter-sage-hub

Shan2L commented Oct 16, 2026

Uh oh!

sayakpaul Oct 16, 2026

Uh oh!

Shan2L Oct 16, 2026

Uh oh!

sayakpaul Oct 16, 2026

Uh oh!

Shan2L Oct 16, 2026 •

edited

Loading

Uh oh!

sayakpaul Oct 17, 2026

Uh oh!

Shan2L Oct 19, 2026

Uh oh!

sayakpaul left a comment

Uh oh!

Shan2L commented Oct 16, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Oct 17, 2026

Uh oh!

sayakpaul left a comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants



		@is_attention
		class TestAiterSageHubFallback:

Conversation

Shan2L commented Oct 16, 2026

What does this PR do?

Results

Tests

Self-review notes

Before submitting

Who can review?

Uh oh!

sayakpaul Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

Shan2L Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

Shan2L Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul Oct 17, 2026

Choose a reason for hiding this comment

Uh oh!

Shan2L Oct 19, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Shan2L commented Oct 16, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Oct 17, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Shan2L Oct 16, 2026 •

edited

Loading