[V1][Attention] Split triton_attn in triton-only and rocm specific backends (#24648)

Signed-off-by: Burkhard Ringlein <[email protected]>

Author: bringlein

vllm/engine/arg_utils.py

@@ -1494,6 +1494,7 @@ def _is_v1_supported_oracle(self, model_config: ModelConfig) -> bool:
             "FLEX_ATTENTION",
             "TREE_ATTN",
             "XFORMERS_VLLM_V1",
+            "ROCM_ATTN_VLLM_V1",
         ]
         if (envs.is_set("VLLM_ATTENTION_BACKEND")
                 and envs.VLLM_ATTENTION_BACKEND not in V1_BACKENDS):

vllm/platforms/interface.py

@@ -67,6 +67,7 @@ class _Backend(enum.Enum):
     FLEX_ATTENTION = enum.auto()
     TREE_ATTN = enum.auto()
     XFORMERS_VLLM_V1 = enum.auto()
+    ROCM_ATTN_VLLM_V1 = enum.auto()
 
 
 class PlatformEnum(enum.Enum):

vllm/platforms/rocm.py

@@ -231,7 +231,17 @@ def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
                 logger.info("Using Flash Attention backend on V1 engine.")
                 return ("vllm.v1.attention.backends."
                         "rocm_aiter_fa.AiterFlashAttentionBackend")
+            elif (envs.VLLM_ROCM_USE_AITER and
+                envs.VLLM_USE_AITER_UNIFIED_ATTENTION) or \
+                    envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION or \
+                        selected_backend == _Backend.ROCM_ATTN_VLLM_V1:
+                # rocm specific backend, with aiter and/or
+                #   triton prefix-prefill
+                logger.info("Using Rocm/Aiter Attention backend on V1 engine.")
+                return ("vllm.v1.attention.backends."
+                        "rocm_attn.RocmAttentionBackend")
             else:
+                # default case, using triton unified attention
                 logger.info("Using Triton Attention backend on V1 engine.")
                 return ("vllm.v1.attention.backends."
                         "triton_attn.TritonAttentionBackend")