From 6049df4310481cd11dc16bc956de2aca70c24899 Mon Sep 17 00:00:00 2001 From: Abdennacer-Badaoui Date: Mon, 10 Aug 2026 15:20:46 +0000 Subject: [PATCH] Use Hub kernel --- .../en/optimization/attention_backends.md | 2 +- src/diffusers/models/attention_dispatch.py | 83 ++++++++----------- src/diffusers/utils/__init__.py | 2 - src/diffusers/utils/import_utils.py | 21 ----- 4 files changed, 34 insertions(+), 74 deletions(-) diff --git a/docs/source/en/optimization/attention_backends.md b/docs/source/en/optimization/attention_backends.md index 79af0bb00685..08df0e75a952 100644 --- a/docs/source/en/optimization/attention_backends.md +++ b/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 | diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 9414c151fd67..465d9232c7c6 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/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: diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index bc193138cc84..bf87525e6c85 100644 --- a/src/diffusers/utils/__init__.py +++ b/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, diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index 639ed59d6393..d2cf394cd9a7 100644 --- a/src/diffusers/utils/import_utils.py +++ b/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