From 4fdd3e0b74ab5bfcf86341a304cc894c09705d5b Mon Sep 17 00:00:00 2001 From: Adrian Date: Fri, 31 Jul 2026 21:30:09 +0200 Subject: [PATCH] [ROCm][ViT] Detect Triton-AMD kernels at their new aiter location (#40289) Signed-off-by: Adrian Kurt Signed-off-by: Adrian Co-authored-by: Claude --- vllm/platforms/rocm.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) 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(