mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-09-12 12:37:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix last commit not working on older pytorch. (#9346)

Browse Source
This commit is contained in:
comfyanonymous
2025-08-14 20:44:02 -07:00
committed by GitHub
parent f0d5d0111f
commit 4e5c230f6a
1 changed files with 13 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
comfy/ops.py
View File

@@ -32,20 +32,21 @@ def scaled_dot_product_attention(q, k, v, *args, **kwargs):
try: try:
if torch.cuda.is_available(): if torch.cuda.is_available():
from torch.nn.attention import SDPBackend, sdpa_kernel from torch.nn.attention import SDPBackend, sdpa_kernel
import inspect
if "set_priority" in inspect.signature(sdpa_kernel).parameters:
SDPA_BACKEND_PRIORITY = [
SDPBackend.FLASH_ATTENTION,
SDPBackend.EFFICIENT_ATTENTION,
SDPBackend.MATH,
]
SDPA_BACKEND_PRIORITY = [ SDPA_BACKEND_PRIORITY.insert(0, SDPBackend.CUDNN_ATTENTION)
SDPBackend.FLASH_ATTENTION,
SDPBackend.EFFICIENT_ATTENTION,
SDPBackend.MATH,
]
SDPA_BACKEND_PRIORITY.insert(0, SDPBackend.CUDNN_ATTENTION) def scaled_dot_product_attention(q, k, v, *args, **kwargs):
with sdpa_kernel(SDPA_BACKEND_PRIORITY, set_priority=True):
def scaled_dot_product_attention(q, k, v, *args, **kwargs): return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
# Use this (rather than the decorator syntax) to eliminate graph else:
# break for pytorch < 2.9 logging.warning("Torch version too old to set sdpa backend priority.")
with sdpa_kernel(SDPA_BACKEND_PRIORITY, set_priority=True):
return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
except (ModuleNotFoundError, TypeError): except (ModuleNotFoundError, TypeError):
logging.warning("Could not set sdpa backend priority.") logging.warning("Could not set sdpa backend priority.")