[diffusion] multi-platform: add Sage Attention 3 Support for sm 120 (RTX5090) (#15382)

Co-authored-by: Mengxi Li <marcyleemx@gmail.com>
This commit is contained in:
ryang
2025-12-19 21:27:02 +08:00
committed by GitHub
co-authored by Mengxi Li
parent 46be74b4b4
commit 1e58248808
7 changed files with 42 additions and 3 deletions
@@ -269,6 +269,13 @@ def is_blackwell():
return torch.cuda.get_device_capability()[0] == 10
@lru_cache(maxsize=1)
def is_sm120():
if not is_cuda():
return False
return torch.cuda.get_device_capability()[0] == 12
@lru_cache(maxsize=1)
def is_hpu() -> bool:
return hasattr(torch, "hpu") and torch.hpu.is_available()