diff --git a/docs/source/en/optimization/attention_backends.md b/docs/source/en/optimization/attention_backends.md index a33fbe89815f..7823d195d773 100644 --- a/docs/source/en/optimization/attention_backends.md +++ b/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 | diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 18e575d8ebbc..9c0213d86875 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/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], diff --git a/tests/models/test_attention_dispatch.py b/tests/models/test_attention_dispatch.py index 478fb2853c0f..db86634ae50f 100644 --- a/tests/models/test_attention_dispatch.py +++ b/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 diff --git a/tests/models/testing_utils/attention.py b/tests/models/testing_utils/attention.py index 4589ef1d1b64..fcb1544dc4a7 100644 --- a/tests/models/testing_utils/attention.py +++ b/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