How to run the latest version of cuDNN with pytorch?

1

The latest version of cuDNN has improvements for FP8 on our SM120 hardware (Blackwell 6000) and we would like to test the performance gains.

The pytorch nightly doesn’t use the latest version of cuDNN, should we manually install cuDNN and then build pytorch from source, or is there a better way?

Thanks!

2

You could try to replace the cuDNN wheel with an updated one via:

# install PyTorch first
# this will also install nvidia-cudnn-cuXX
...
# uninstall the already installed cuDNN wheel
pip uninstall nvidia-cudnn-cuXX
...
# install the one you want
pip install nvidia-cudnn-cuXX==X.Y.Z.W

Depending on the changes you would need to check if these cuDNN kernels are guarded by a build-time check, which would ignore the runtime version. In this case you would need to rebuild PyTorch from source.