[ROCm][ViT] Detect Triton-AMD kernels at their new aiter location (#40289)
Signed-off-by: Adrian Kurt <info@zzit.ch> Signed-off-by: Adrian <info@zipfuzapfu.ch> Co-authored-by: Claude <noreply@anthropic.com>
Adrian committed
Jul 31, 2026 at 21:30 UTC
4fdd3e0b74ab5bfcf86341a304cc894c09705d5b
1 file changed
+15
-3
vllm/platforms/rocm.py
+15
-3
@@ -413,9 +413,21 @@ def flash_attn_triton_available() -> bool:
413
try:
414
from importlib.util import find_spec
415
416
- if find_spec("flash_attn") is None:
417
- return False
418
- if find_spec("flash_attn.flash_attn_triton_amd") is None:
416
+ # Locate the Triton-AMD kernels. Older ROCm/flash-attention (pre
417
+ # 2026-03) shipped them as the flash_attn.flash_attn_triton_amd
418
+ # subpackage. The main_perf migration commit 3f94643 moved them
419
+ # into aiter at aiter.ops.triton._triton_kernels.flash_attn_triton_amd,
420
+ # so accept either location.
421
+ def _has_spec(name: str) -> bool:
422
+ try:
423
+ return find_spec(name) is not None
424
+ except (ImportError, ValueError):
425
+ return False
426
+
427
+ if not (
428
+ _has_spec("flash_attn.flash_attn_triton_amd")
429
+ or _has_spec("aiter.ops.triton._triton_kernels.flash_attn_triton_amd")
430
+ ):
431
return False
432
if os.environ.get("FLASH_ATTENTION_TRITON_AMD_ENABLE") != "TRUE":
433
logger.info_once(