[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:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user