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

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

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