Add aiter_sage_hub attention backend for AMD RDNA4 (gfx1201)

docs/source/en/optimization/attention_backends.md

@@ -157,6 +157,7 @@ Refer to the table below for a complete list of available attention backends and
 | `flash_varlen` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | Variable length FlashAttention |
 | `flash_varlen_hub` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | Variable length FlashAttention from kernels |
 | `aiter_fa2_hub` | [AI Tensor Engine for ROCm](https://github.com/ROCm/aiter) | FlashAttention-2 for AMD ROCm from kernels |
+| `aiter_sage_hub` | [AI Tensor Engine for ROCm](https://github.com/ROCm/aiter) | Quantized attention (INT8 QK, FP8 PV) for AMD RDNA4 (gfx1201) from kernels |
 | `flash_4_hub` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-4 |
 | `_flash_3` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-3 |
 | `_flash_varlen_3` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | Variable length FlashAttention-3 |

src/diffusers/models/attention_dispatch.py

@@ -231,6 +231,8 @@ class AttentionBackendName(str, Enum):
 
     # `aiter` (via the `kernels-community/aiter-flash-attn-ck` Hub kernel)
     AITER_FA2_HUB = "aiter_fa2_hub"
+    # `aiter` SageAttention for gfx1201 (via the `kernels-community/aiter-sage-attention` Hub kernel)
+    AITER_SAGE_HUB = "aiter_sage_hub"
 
     # PyTorch native
     FLEX = "flex"
@@ -373,6 +375,11 @@ class _HubKernelConfig:
         function_attr="flash_attn_func",
         version=1,
     ),
+    AttentionBackendName.AITER_SAGE_HUB: _HubKernelConfig(
+        repo_id="kernels-community/aiter-sage-attention",
+        function_attr="sageattn",
+        version=1,
+    ),
 }
 
 
@@ -554,6 +561,7 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None
         AttentionBackendName.SAGE_BLACKWELL_HUB,
         AttentionBackendName.FLASH_4_HUB,
         AttentionBackendName.AITER_FA2_HUB,
+        AttentionBackendName.AITER_SAGE_HUB,
     ]:
         if not is_kernels_available():
             raise RuntimeError(
@@ -3604,6 +3612,106 @@ def _aiter_flash_attention(
     return out
 
 
+def _aiter_sage_attention_hub_supported(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: torch.Tensor | None,
+    dropout_p: float,
+    is_causal: bool,
+    enable_gqa: bool,
+    return_lse: bool,
+) -> bool:
+    return (
+        query.shape == key.shape == value.shape
+        and query.shape[-1] == 128
+        and query.dtype == key.dtype == value.dtype == torch.bfloat16
+        and attn_mask is None
+        and dropout_p == 0.0
+        and not is_causal
+        and not enable_gqa
+        and not return_lse
+    )
+
+
+def _aiter_sage_attention_hub_forward_op(
+    ctx: torch.autograd.function.FunctionCtx,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: torch.Tensor | None = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: float | None = None,
+    enable_gqa: bool = False,
+    return_lse: bool = False,
+    _save_ctx: bool = True,
+    _parallel_config: "ParallelConfig" | None = None,
+):
+    if not _aiter_sage_attention_hub_supported(
+        query, key, value, attn_mask, dropout_p, is_causal, enable_gqa, return_lse
+    ):
+        return _native_attention_forward_op(
+            ctx,
+            query,
+            key,
+            value,
+            attn_mask,
+            dropout_p,
+            is_causal,
+            scale,
+            enable_gqa,
+            return_lse,
+            _save_ctx,
+            _parallel_config,
+        )
+    func = _HUB_KERNELS_REGISTRY[AttentionBackendName.AITER_SAGE_HUB].kernel_fn
+    return func(q=query, k=key, v=value, tensor_layout="NHD", is_causal=is_causal, sm_scale=scale)
+
+
+@_AttentionBackendRegistry.register(
+    AttentionBackendName.AITER_SAGE_HUB,
+    # The gfx1201 kernel takes bf16 Q/K/V of equal shape with head dim 128, non-causal and unmasked. Other calls (for
+    # example head-dim-64 VAE attention or cross-attention that follow the active backend) fall back to native SDPA.
+    constraints=[_check_device_cuda, _check_shape],
+    supports_context_parallel=True,
+)
+def _aiter_sage_attention_hub(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: torch.Tensor | None = None,
+    is_causal: bool = False,
+    scale: float | None = None,
+    return_lse: bool = False,
+    _parallel_config: "ParallelConfig" | None = None,
+) -> torch.Tensor:
+    if _parallel_config is None:
+        if not _aiter_sage_attention_hub_supported(query, key, value, attn_mask, 0.0, is_causal, False, return_lse):
+            return _native_attention(query, key, value, attn_mask, 0.0, is_causal, scale, False, return_lse)
+        return _aiter_sage_attention_hub_forward_op(
+            None, query, key, value, attn_mask, 0.0, is_causal, scale, False, return_lse
+        )
+    if _parallel_config.context_parallel_config.ring_degree > 1:
+        raise NotImplementedError(
+            "The AITER_SAGE_HUB backend does not support `ring_degree > 1`, since the kernel does not return the LSE."
+        )
+    return _templated_context_parallel_attention(
+        query,
+        key,
+        value,
+        attn_mask,
+        0.0,
+        is_causal,
+        scale,
+        False,
+        return_lse,
+        forward_op=_aiter_sage_attention_hub_forward_op,
+        backward_op=_sage_attention_backward_op,
+        _parallel_config=_parallel_config,
+    )
+
+
 @_AttentionBackendRegistry.register(
     AttentionBackendName.FLEX,
     constraints=[_check_attn_mask_or_causal, _check_device, _check_shape],

tests/models/test_attention_dispatch.py

@@ -23,8 +23,13 @@
 import torch.nn.functional as F
 
 from diffusers.models._modeling_parallel import ContextParallelConfig, ParallelConfig
+from diffusers.models.attention_dispatch import (
+    _HUB_KERNELS_REGISTRY,
+    AttentionBackendName,
+    _AttentionBackendRegistry,
+    dispatch_attention_fn,
+)
 from diffusers.models.attention_dispatch import attention_backend as attention_backend_ctx
-from diffusers.models.attention_dispatch import dispatch_attention_fn
 
 from ..testing_utils import (
     is_attention,
@@ -35,6 +40,7 @@
     require_torch_multi_accelerator,
     torch_device,
 )
+from .testing_utils.attention import _IS_GFX1201
 from .testing_utils.parallelism import DEVICE_CONFIG, _find_free_port
 
 
@@ -231,3 +237,78 @@ def make_qkv(seq_len):
                         )
             finally:
                 torch.compiler.reset()
+
+
+@is_attention
+class TestAiterSageHubFallback:
+    """`aiter_sage_hub` runs the gfx1201 kernel only on inputs it supports and falls back to native SDPA otherwise.
+
+    The kernel is replaced by a stub, so this runs without a gfx1201 GPU or the Hub kernel.
+    """
+
+    @pytest.fixture
+    def kernel_calls(self, monkeypatch):
+        calls = []
+
+        def fake_sageattn(q, k, v, tensor_layout, is_causal, sm_scale):
+            calls.append(tensor_layout)
+            return torch.zeros_like(q)
+
+        config = _HUB_KERNELS_REGISTRY[AttentionBackendName.AITER_SAGE_HUB]
+        monkeypatch.setattr(config, "kernel_fn", fake_sageattn)
+        return calls
+
+    @staticmethod
+    def _run(query, key, value, **kwargs):
+        backend_fn = _AttentionBackendRegistry._backends[AttentionBackendName.AITER_SAGE_HUB]
+        return backend_fn(query, key, value, **kwargs)
+
+    @staticmethod
+    def _qkv(seq_q=16, seq_kv=16, head_dim=128, dtype=torch.bfloat16):
+        q = torch.randn(1, seq_q, 2, head_dim, dtype=dtype)
+        k = torch.randn(1, seq_kv, 2, head_dim, dtype=dtype)
+        v = torch.randn(1, seq_kv, 2, head_dim, dtype=dtype)
+        return q, k, v
+
+    def test_supported_inputs_use_kernel(self, kernel_calls):
+        self._run(*self._qkv())
+        assert len(kernel_calls) == 1
+
+    @pytest.mark.parametrize(
+        "qkv_kwargs, call_kwargs",
+        [
+            ({"head_dim": 64}, {}),
+            ({"dtype": torch.float32}, {}),
+            ({"seq_kv": 8}, {}),
+            ({}, {"is_causal": True}),
+            ({}, {"attn_mask": torch.ones(16, 16, dtype=torch.bool)}),
+        ],
+        ids=["head_dim_64", "fp32", "cross_attention", "causal", "mask"],
+    )
+    def test_unsupported_inputs_fall_back_to_native(self, kernel_calls, qkv_kwargs, call_kwargs):
+        query, key, value = self._qkv(**qkv_kwargs)
+        out = self._run(query, key, value, **call_kwargs)
+        ref = _AttentionBackendRegistry._backends[AttentionBackendName.NATIVE](query, key, value, **call_kwargs)
+        assert kernel_calls == []
+        assert torch.equal(out, ref)
+
+    def test_ring_attention_raises(self, kernel_calls):
+        config = ParallelConfig(context_parallel_config=ContextParallelConfig(ring_degree=2))
+        with pytest.raises(NotImplementedError, match="ring_degree"):
+            self._run(*self._qkv(), _parallel_config=config)
+
+
+@is_attention
+@pytest.mark.skipif(not _IS_GFX1201, reason="A gfx1201 GPU is required for the aiter_sage_hub kernel.")
+@pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available.")
+class TestAiterSageHubKernel:
+    def test_output_close_to_native(self):
+        torch.manual_seed(0)
+        query, key, value = (torch.randn(1, 1024, 4, 128, device=torch_device, dtype=torch.bfloat16) for _ in range(3))
+        ref = dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE)
+        with attention_backend_ctx(AttentionBackendName.AITER_SAGE_HUB):
+            out = dispatch_attention_fn(query, key, value)
+        # The quantized kernel never matches SDPA exactly, so equality would mean the fallback ran.
+        assert not torch.equal(out, ref)
+        cos = F.cosine_similarity(out.float().flatten(), ref.float().flatten(), dim=0)
+        assert cos > 0.99

tests/models/testing_utils/attention.py

@@ -41,6 +41,12 @@
 # kernel only loads on SM120 (consumer/workstation Blackwell). `a` targets are architecture
 # specific, so neither SM100 nor SM121 is covered.
 _IS_SM120 = _CUDA_AVAILABLE and torch.cuda.get_device_capability() == (12, 0)
+# `kernels-community/aiter-sage-attention` is built for gfx1201 (RDNA4) only.
+_IS_GFX1201 = (
+    _CUDA_AVAILABLE
+    and torch.version.hip is not None
+    and torch.cuda.get_device_properties(0).gcnArchName.split(":")[0] == "gfx1201"
+)
 
 _PARAM_NATIVE_CUDNN = pytest.param(
     AttentionBackendName._NATIVE_CUDNN,
@@ -125,6 +131,18 @@
     ],
 )
 
+_PARAM_AITER_SAGE_HUB = pytest.param(
+    AttentionBackendName.AITER_SAGE_HUB,
+    id="aiter_sage_hub",
+    marks=[
+        pytest.mark.skipif(not _IS_GFX1201, reason="A gfx1201 GPU is required for the aiter_sage_hub backend."),
+        pytest.mark.skipif(
+            not is_kernels_available(),
+            reason="`kernels` package is required for aiter_sage_hub backend. Install with `pip install kernels`.",
+        ),
+    ],
+)
+
 # All backends under test.
 _ALL_BACKEND_PARAMS = [
     _PARAM_NATIVE_CUDNN,
@@ -134,6 +152,7 @@
     _PARAM_FLASH_3_VARLEN_HUB,
     _PARAM_SAGE_HUB,
     _PARAM_SAGE_BLACKWELL_HUB,
+    _PARAM_AITER_SAGE_HUB,
 ]
 
 # Backends that perform non-deterministic operations and therefore cannot run when