File size: 842 Bytes
d0bfdd6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
from contextlib import contextmanager
import torch
try:
from flash_attn import flash_attn_varlen_qkvpacked_func as flash_varlen_qkvpacked_attn
except ImportError:
flash_varlen_qkvpacked_attn = None
try:
from sageattention import sageattn as sage_attn
except ImportError:
sage_attn = None
try:
from comfy.ldm.modules.attention import comfy_optimized_attention as comfy_attn
except ImportError:
comfy_attn = None
from torch.nn.attention import SDPBackend, sdpa_kernel
backends = []
if torch.cuda.get_device_properties(0).major < 7:
backends.append(SDPBackend.MATH)
if torch.cuda.get_device_properties(0).major >= 9.0:
backends.append(SDPBackend.CUDNN_ATTENTION)
else:
backends.append(SDPBackend.EFFICIENT_ATTENTION)
@contextmanager
def sdpa_attn_ctx():
with sdpa_kernel(backends):
yield
|