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