From 2220d111c8b036f094eb465400fdf962626e4afa Mon Sep 17 00:00:00 2001 From: Alex Harper Date: Wed, 12 Aug 2026 21:40:11 -0400 Subject: [PATCH] Query pytorch for aotriton support instead of listing its lib directory (#15412) --- comfy/model_management.py | 36 ++++++++++++++++++++++-------------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 65599424b..15c03dc77 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -490,28 +490,36 @@ try: except: rocm_version = (6, -1) - def aotriton_supported(gpu_arch): - path = torch.__path__[0] - path = os.path.join(os.path.join(path, "lib"), "aotriton.images") - gfx = set(map(lambda a: a[4:], filter(lambda a: a.startswith("amd-gfx"), os.listdir(path)))) - if gpu_arch in gfx: - return True - if "{}x".format(gpu_arch[:-1]) in gfx: - return True - if "{}xx".format(gpu_arch[:-2]) in gfx: - return True - return False + def aotriton_supported(): + """Whether pytorch reports flash attention as usable on this gpu. + + can_use_flash_attention() evaluates runtime eligibility for the given + parameters; on a ROCm build that includes checking the gpu arch against the + kernel images AOTriton was compiled for. Querying it avoids assuming where + those images live inside the torch install. The probe tensor is shaped and + typed to pass the unrelated SDPA checks, so False means no hardware support + rather than a rejected shape. + """ + try: + if not torch.backends.cuda.is_flash_attention_available(): # not built with flash attention + return False + q = torch.empty((1, 1, 8, 64), dtype=torch.float16, device=get_torch_device()) + params = torch.backends.cuda.SDPAParams(q, q, q, None, 0.0, False, False) + return torch.backends.cuda.can_use_flash_attention(params, False) + except (AttributeError, RuntimeError, TypeError) as e: + logging.warning("Could not query aotriton support: {}".format(e)) + return False logging.info("AMD arch: {}".format(arch)) logging.info("ROCm version: {}".format(rocm_version)) if args.use_split_cross_attention == False and args.use_quad_cross_attention == False: - if aotriton_supported(arch): # AMD efficient attention implementation depends on aotriton. + if aotriton_supported(): # AMD efficient attention implementation depends on aotriton. if torch_version_numeric >= (2, 7): # works on 2.6 but doesn't actually seem to improve much if any((a in arch) for a in ["gfx90a", "gfx942", "gfx950", "gfx1100", "gfx1101", "gfx1150", "gfx1151"]): # TODO: more arches, TODO: gfx950 ENABLE_PYTORCH_ATTENTION = True if rocm_version >= (7, 0): - if any((a in arch) for a in ["gfx1200", "gfx1201"]): - ENABLE_PYTORCH_ATTENTION = True + if any((a in arch) for a in ["gfx1200", "gfx1201"]): + ENABLE_PYTORCH_ATTENTION = True if torch_version_numeric >= (2, 7) and rocm_version >= (6, 4): if any((a in arch) for a in ["gfx1200", "gfx1201", "gfx950"]): # TODO: more arches, "gfx942" gives error on pytorch nightly 2.10 1013 rocm7.0 SUPPORT_FP8_OPS = True