diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index e39b279124..d7e5bf02f2 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -413,9 +413,21 @@ def flash_attn_triton_available() -> bool: try: from importlib.util import find_spec - if find_spec("flash_attn") is None: - return False - if find_spec("flash_attn.flash_attn_triton_amd") is None: + # Locate the Triton-AMD kernels. Older ROCm/flash-attention (pre + # 2026-03) shipped them as the flash_attn.flash_attn_triton_amd + # subpackage. The main_perf migration commit 3f94643 moved them + # into aiter at aiter.ops.triton._triton_kernels.flash_attn_triton_amd, + # so accept either location. + def _has_spec(name: str) -> bool: + try: + return find_spec(name) is not None + except (ImportError, ValueError): + return False + + if not ( + _has_spec("flash_attn.flash_attn_triton_amd") + or _has_spec("aiter.ops.triton._triton_kernels.flash_attn_triton_amd") + ): return False if os.environ.get("FLASH_ATTENTION_TRITON_AMD_ENABLE") != "TRUE": logger.info_once(