[FA2] Use kernels-community/aiter-flash-attn-ck Hub kernel for ROCm FlashAttention (drop aiter dependency)

docs/source/en/optimization/attention_backends.md

@@ -156,7 +156,7 @@ Refer to the table below for a complete list of available attention backends and
 | `flash_hub` | [FlashAttention](https://github.com/Dao-AILab/flash-attention) | FlashAttention-2 from kernels |
 | `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` | [AI Tensor Engine for ROCm](https://github.com/ROCm/aiter) | FlashAttention for AMD ROCm |
+| `aiter_fa2_hub` | [AI Tensor Engine for ROCm](https://github.com/ROCm/aiter) | FlashAttention-2 for AMD ROCm 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

@@ -32,8 +32,6 @@
 
 from ..utils import (
     get_logger,
-    is_aiter_available,
-    is_aiter_version,
     is_flash_attn_3_available,
     is_flash_attn_available,
     is_flash_attn_version,
@@ -57,7 +55,6 @@
     from ._modeling_parallel import ParallelConfig
 
 _REQUIRED_FLASH_VERSION = "2.6.3"
-_REQUIRED_AITER_VERSION = "0.1.5"
 _REQUIRED_SAGE_VERSION = "2.1.1"
 _REQUIRED_FLEX_VERSION = "2.5.0"
 _REQUIRED_XLA_VERSION = "2.2"
@@ -67,7 +64,6 @@
 
 _CAN_USE_FLASH_ATTN = is_flash_attn_available() and is_flash_attn_version(">=", _REQUIRED_FLASH_VERSION)
 _CAN_USE_FLASH_ATTN_3 = is_flash_attn_3_available()
-_CAN_USE_AITER_ATTN = is_aiter_available() and is_aiter_version(">=", _REQUIRED_AITER_VERSION)
 _CAN_USE_SAGE_ATTN = is_sageattention_available() and is_sageattention_version(">=", _REQUIRED_SAGE_VERSION)
 _CAN_USE_FLEX_ATTN = is_torch_version(">=", _REQUIRED_FLEX_VERSION)
 _CAN_USE_NPU_ATTN = is_torch_npu_available()
@@ -108,16 +104,6 @@
     flash_attn_3_func = None
     flash_attn_3_varlen_func = None
 
-if _CAN_USE_AITER_ATTN:
-    try:
-        from aiter import flash_attn_func as aiter_flash_attn_func
-    except (ImportError, OSError, RuntimeError) as e:
-        logger.warning(f"aiter failed to import: {e}. Falling back to native attention.")
-        _CAN_USE_AITER_ATTN = False
-        aiter_flash_attn_func = None
-else:
-    aiter_flash_attn_func = None
-
 if _CAN_USE_SAGE_ATTN:
     try:
         from sageattention import (
@@ -235,8 +221,8 @@ class AttentionBackendName(str, Enum):
     _FLASH_3_HUB = "_flash_3_hub"
     _FLASH_3_VARLEN_HUB = "_flash_3_varlen_hub"
 
-    # `aiter`
-    AITER = "aiter"
+    # `aiter` (via the `kernels-community/aiter-flash-attn-ck` Hub kernel)
+    AITER_FA2_HUB = "aiter_fa2_hub"
 
     # PyTorch native
     FLEX = "flex"
@@ -368,6 +354,11 @@ class _HubKernelConfig:
         function_attr="flash_attn_func",
         version=0,
     ),
+    AttentionBackendName.AITER_FA2_HUB: _HubKernelConfig(
+        repo_id="kernels-community/aiter-flash-attn-ck",
+        function_attr="flash_attn_func",
+        version=1,
+    ),
 }
 
 
@@ -490,6 +481,12 @@ def _check_qkv_dtype_bf16_or_fp16(query: torch.Tensor, key: torch.Tensor, value:
         raise ValueError("Query, key, and value must be either bfloat16 or float16.")
 
 
+def _check_qkv_dtype_bf16(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs) -> None:
+    _check_qkv_dtype_match(query, key, value)
+    if query.dtype != torch.bfloat16:
+        raise ValueError("Query, key, and value must be bfloat16.")
+
+
 def _check_shape(
     query: torch.Tensor,
     key: torch.Tensor,
@@ -534,6 +531,7 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None
         AttentionBackendName._FLASH_3_VARLEN_HUB,
         AttentionBackendName.SAGE_HUB,
         AttentionBackendName.FLASH_4_HUB,
+        AttentionBackendName.AITER_FA2_HUB,
     ]:
         if not is_kernels_available():
             raise RuntimeError(
@@ -549,12 +547,6 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None
                 f"Backend '{backend.value}' needs to be used with a `kernels` version of at least 0.12.3. Please update with `pip install -U kernels`."
             )
 
-    elif backend == AttentionBackendName.AITER:
-        if not _CAN_USE_AITER_ATTN:
-            raise RuntimeError(
-                f"Aiter Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `aiter>={_REQUIRED_AITER_VERSION}`."
-            )
-
     elif backend in [
         AttentionBackendName.SAGE,
         AttentionBackendName.SAGE_VARLEN,
@@ -3371,8 +3363,10 @@ def _flash_varlen_attention_3(
 
 
 @_AttentionBackendRegistry.register(
-    AttentionBackendName.AITER,
-    constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16, _check_shape],
+    AttentionBackendName.AITER_FA2_HUB,
+    # The `kernels-community/aiter-flash-attn-ck` CK kernel only ships bf16 `mha_fwd` instances;
+    # fp16 raises a cryptic "invalid argument for fmha_fwd" from CK, so reject it up front.
+    constraints=[_check_device_cuda, _check_qkv_dtype_bf16, _check_shape],
 )
 def _aiter_flash_attention(
     query: torch.Tensor,
@@ -3388,31 +3382,20 @@ def _aiter_flash_attention(
     if attn_mask is not None:
         raise ValueError("`attn_mask` is not supported for aiter attention")
 
-    if not return_lse and torch.is_grad_enabled():
-        # aiter requires return_lse=True by assertion when gradients are enabled.
-        out, lse, *_ = aiter_flash_attn_func(
-            q=query,
-            k=key,
-            v=value,
-            dropout_p=dropout_p,
-            softmax_scale=scale,
-            causal=is_causal,
-            return_lse=True,
-        )
-    else:
-        out = aiter_flash_attn_func(
-            q=query,
-            k=key,
-            v=value,
-            dropout_p=dropout_p,
-            softmax_scale=scale,
-            causal=is_causal,
-            return_lse=return_lse,
-        )
-        if return_lse:
-            out, lse, *_ = out
-
-    return (out, lse) if return_lse else out
+    func = _HUB_KERNELS_REGISTRY[AttentionBackendName.AITER_FA2_HUB].kernel_fn
+    out = func(
+        q=query,
+        k=key,
+        v=value,
+        dropout_p=dropout_p,
+        softmax_scale=scale,
+        causal=is_causal,
+        return_lse=return_lse,
+    )
+    if return_lse:
+        out, lse, *_ = out
+        return out, lse
+    return out
 
 
 @_AttentionBackendRegistry.register(
@@ -3684,7 +3667,7 @@ def _native_flash_attention(
     _parallel_config: "ParallelConfig" | None = None,
 ) -> torch.Tensor:
     if attn_mask is not None:
-        raise ValueError("`attn_mask` is not supported for aiter attention")
+        raise ValueError("`attn_mask` is not supported for native flash attention")
 
     lse = None
     if _parallel_config is None and not return_lse:

src/diffusers/utils/__init__.py

@@ -65,8 +65,6 @@
     get_objects_from_module,
     is_accelerate_available,
     is_accelerate_version,
-    is_aiter_available,
-    is_aiter_version,
     is_auto_round_available,
     is_av_available,
     is_better_profanity_available,

src/diffusers/utils/import_utils.py

@@ -214,7 +214,6 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[b
 _sageattention_available, _sageattention_version = _is_package_available("sageattention")
 _flash_attn_available, _flash_attn_version = _is_package_available("flash_attn")
 _flash_attn_3_available, _flash_attn_3_version = _is_package_available("flash_attn_3")
-_aiter_available, _aiter_version = _is_package_available("aiter", get_dist_name=True)
 _kornia_available, _kornia_version = _is_package_available("kornia")
 _nvidia_modelopt_available, _nvidia_modelopt_version = _is_package_available("modelopt", get_dist_name=True)
 _auto_round_available, _auto_round_version = _is_package_available("auto_round")
@@ -415,10 +414,6 @@ def is_flash_attn_3_available():
     return _flash_attn_3_available
 
 
-def is_aiter_available():
-    return _aiter_available
-
-
 def is_kornia_available():
     return _kornia_available
 
@@ -941,22 +936,6 @@ def is_flash_attn_version(operation: str, version: str):
     return compare_versions(parse(_flash_attn_version), operation, version)
 
 
-@cache
-def is_aiter_version(operation: str, version: str):
-    """
-    Compares the current aiter version to a given reference with an operation.
-
-    Args:
-        operation (`str`):
-            A string representation of an operator, such as `">"` or `"<="`
-        version (`str`):
-            A version string
-    """
-    if not _aiter_available:
-        return False
-    return compare_versions(parse(_aiter_version), operation, version)
-
-
 def get_objects_from_module(module):
     """
     Returns a dict of object names and values in a module, while skipping private/internal objects