From 46e96e18e37b97b596de36f04d5c8104a801e8ce Mon Sep 17 00:00:00 2001 From: TJian Date: Tue, 27 Jan 2026 20:14:49 +0800 Subject: [PATCH] [ROCm] [Perf] Add AITER Flash Attention (#941) Signed-off-by: tjtanaa --- .../attention/backends/flash_attn.py | 9 ++++++-- .../diffusion/attention/backends/utils/fa.py | 14 ++++++++++++ vllm_omni/diffusion/attention/selector.py | 22 +++++++++++++++++++ 3 files changed, 43 insertions(+), 2 deletions(-) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 7fad4ce26..e9102b1d3 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func from vllm.logger import init_logger from vllm_omni.diffusion.attention.backends.abstract import ( @@ -10,7 +9,13 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionImpl, AttentionMetadata, ) -from vllm_omni.diffusion.attention.backends.utils.fa import _pad_input, _unpad_input, _upad_input +from vllm_omni.diffusion.attention.backends.utils.fa import ( + _pad_input, + _unpad_input, + _upad_input, + flash_attn_func, + flash_attn_varlen_func, +) logger = init_logger(__name__) diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index d89082e71..6a02435a7 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -15,6 +15,20 @@ import torch import torch.nn.functional as F +from vllm_omni.utils.platform_utils import is_rocm + +if is_rocm(): + from vllm._aiter_ops import is_aiter_found_and_supported + + # Choose to enable this by default on ROCm + # Whenever possible as it is the fastest backend + if is_aiter_found_and_supported(): + from aiter import flash_attn_func, flash_attn_varlen_func # noqa: F401 + else: + raise ImportError("Aiter is not found and supported on currentROCm device.") +else: + from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401 + def _index_first_axis(tensor, indices): """ diff --git a/vllm_omni/diffusion/attention/selector.py b/vllm_omni/diffusion/attention/selector.py index 706e79e7b..ed5946098 100644 --- a/vllm_omni/diffusion/attention/selector.py +++ b/vllm_omni/diffusion/attention/selector.py @@ -81,6 +81,28 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]: "Falling back to TORCH_SDPA backend.""" ) backend_name = "TORCH_SDPA" + elif detect_device_type() == "cuda" and is_rocm(): + from vllm._aiter_ops import is_aiter_found_and_supported + + compute_capability = torch.cuda.get_device_capability() + major, minor = compute_capability + + # Choose to enable this by default on ROCm + # Whenever possible as it is the fastest backend + # is_aiter_found_and_supported() checks if aiter library is found + # and is aiter supported on the current platform + # aiter currently only is supported on gfx942 and gfx950 + # https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py + if is_aiter_found_and_supported() and 90 < major * 10 + minor < 100: + if backend_name is None: + backend_name = "FLASH_ATTN" + else: + if backend_name == "FLASH_ATTN": + logger.warning( + "Flash Attention requires `aiter` library which is only supported " + "on gfx942 and gfx950. Falling back to TORCH_SDPA backend." + ) + backend_name = "TORCH_SDPA" if backend_name is not None: backend_name_upper = backend_name.upper()