[ROCm] Enable AITER backend with ring attention (#3511)

Signed-off-by: Aleksi Vesanto <avesanto@amd.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/attention/parallel/ring.py

@@ -10,7 +10,7 @@ import torch
 from vllm.logger import init_logger
 
 # import torch.distributed as dist # Not used directly here, but good practice if needed
-from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_FA3, HAS_FLASH_ATTN
+from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_AITER, HAS_FA3, HAS_FLASH_ATTN
 from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
 from vllm_omni.diffusion.attention.parallel.base import (
     ParallelAttentionContext,
@@ -110,13 +110,12 @@ class RingParallelAttention:
         if backend_pref is not None:
             backend_pref = backend_pref.lower()
 
-        # Determine attention type with fallback chain: FA3 -> FA2 -> SDPA
         # FP32 is not supported by Flash Attention, force SDPA
         if query.dtype == torch.float32:
             backend_pref = "sdpa"
-        elif not HAS_FA3 and not HAS_FLASH_ATTN:
+        elif not HAS_FA3 and not HAS_FLASH_ATTN and not HAS_AITER:
             if backend_pref != "sdpa":
-                logger.warning_once("Flash Attention (FA2/FA3) is not available! Force enabling SDPA.")
+                logger.warning_once("Flash Attention (FA2/FA3/AITER) is not available! Force enabling SDPA.")
             backend_pref = "sdpa"
 
         # Extract joint tensors
@@ -147,7 +146,13 @@ class RingParallelAttention:
         from vllm_omni.diffusion.attention.backends.ring_flash_attn import ring_flash_attn_func
 
         # Prefer FA3 over FA2 for better performance (FA3 supports Ampere/Ada/Hopper)
-        attn_type = AttnType.FA3 if HAS_FA3 else AttnType.FA
+        # On ROCm, use AITER
+        if HAS_FA3:
+            attn_type = AttnType.FA3
+        elif HAS_AITER:
+            attn_type = AttnType.AITER
+        else:
+            attn_type = AttnType.FA
 
         return ring_flash_attn_func(
             query,