RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Query pytorch for aotriton support instead of listing its lib directory (#15412)

This commit is contained in:
Alex Harper 2026-08-12 21:40:11 -04:00 committed by GitHub
parent 6b30dc2068
commit 2220d111c8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 22 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
comfy/model_management.py
View File

@ -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