restore fa

ROCm · Nov 20, 2024 · e5cf3da · e5cf3da
1 parent 1348c23
commit e5cf3da
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/vllm/attention/backends/rocm_flash_attn.py b/vllm/attention/backends/rocm_flash_attn.py
@@ -3,7 +3,6 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
 
 import torch
-from flash_attn import flash_attn_varlen_func  # noqa: F401
 
 import vllm.envs as envs
 from vllm import _custom_ops as ops
@@ -508,8 +507,12 @@ def __init__(
             if not current_platform.has_device_capability(90):
                 self.use_naive_attn = True
             else:
-                self.attn_func = flash_attn_varlen_func
-                logger.debug("Using CK FA in ROCmBackend")
+                try:
+                    from flash_attn import flash_attn_varlen_func  # noqa: F401
+                    self.attn_func = flash_attn_varlen_func
+                    logger.debug("Using CK FA in ROCmBackend")
+                except ModuleNotFoundError:
+                    self.use_naive_attn = True
 
             if self.use_naive_attn:
                 if logits_soft_cap is not None: