Add support for Nvidia Blackwell (sm_120)
Description:
In PKGBUILD script (https://gitlab.archlinux.org/archlinux/packaging/packages/python-pytorch/-/blob/main/PKGBUILD?ref_type=heads#L232), the cuda capabilities are limited from sm61 to sm90a. Nvidia Blackwell series GPUs are not compatible with them. Typical runtime error message would looks like:
NVIDIA GeForce RTX 5090 with CUDA capability sm_120 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_61 sm_70 sm_72 sm_75 sm_80 sm_86 sm_87 sm_89 sm_90 sm_90a.
RuntimeError: CUDA error: no kernel image is available for execution on the device
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
It would be RTX5xxx friendly to append 12.0 to the TORCH_CUDA_ARCH_LIST to fix this problem (confirmed on my 5090).
Additional info:
- package version(s):
2.7.0
Steps to reproduce:
- Inference a small Pytorch model, such as real-esrgan
- And error message will pop out