probe-attention-backend.py
674 Bytes
import os, sys, torch
from diffsynth.core.attention import attention_forward
A = sys.modules[attention_forward.__module__]
print("DIFFSYNTH_ATTENTION_IMPLEMENTATION env =", os.environ.get("DIFFSYNTH_ATTENTION_IMPLEMENTATION"))
print("selected =", A.ATTENTION_IMPLEMENTATION,
"| xformers=", A.XFORMERS_AVAILABLE, "flash2=", A.FLASH_ATTN_2_AVAILABLE,
"flash3=", A.FLASH_ATTN_3_AVAILABLE, "flash4=", A.FLASH_ATTN_4_AVAILABLE, "sage=", A.SAGE_ATTN_AVAILABLE)
dev="cuda:1"
q=torch.randn(1,56,14100,128,device=dev,dtype=torch.float16)
o=A.attention_forward(q,q,q)
print("[OK] attention_forward fp16 S=14100:", tuple(o.shape), "finite=", bool(torch.isfinite(o).all()))