[ROCm] [Perf] Add AITER Flash Attention (#941)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

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__)
 

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):
     """

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()