switching to the implementation of MHA in FA

vllm/attention/layer.py

@@ -416,11 +416,14 @@ class MultiHeadAttention(nn.Module):
             backend = _Backend.FLASH_ATTN
             use_upstream_fa = True
 
-        if current_platform.is_rocm() or current_platform.is_xpu():
+        if current_platform.is_xpu():
             # currently, only torch_sdpa is supported on rocm/xpu
             self.attn_backend = _Backend.TORCH_SDPA
+        elif current_platform.is_rocm():
+            self.attn_backend = backend if backend in {
+                _Backend.FLASH_ATTN,
+            } else _Backend.TORCH_SDPA
         else:
-
             self.attn_backend = backend if backend in {
                 _Backend.TORCH_SDPA,
                 _Backend.XFORMERS,
@@ -437,6 +440,10 @@ class MultiHeadAttention(nn.Module):
             if use_upstream_fa:
                 from flash_attn import flash_attn_varlen_func
                 self._flash_attn_varlen_func = flash_attn_varlen_func
+            else:
+                if current_platform.is_rocm():
+                    from flash_attn import flash_attn_varlen_func
+                    self._flash_attn_varlen_func = flash_attn_varlen_func
                 else:
                     from vllm.vllm_flash_attn import flash_attn_varlen_func
                     self._flash_attn_varlen_func = flash_attn_varlen_func