google-deepmind / google-deepmind/mujoco_playground
JAX 0.9.x CUDA: "No visible GPU devices" on RTX 5090 (Blackwell)
- Dominant language
- Python
- Stars
- 2.2k
- Forks
- 359
- Avg merge
- 7d 3h
- Merged PRs (30d)
- 1
Description
JAX CUDA backend fails to detect GPU on RTX 5090 (Blackwell).
**System**:
GPU: NVIDIA GeForce RTX 5090
Driver: 590.48.01
CUDA (from nvidia-smi): 13.1
OS: Ubuntu 24.04
Python: 3.12
**JAX**:
jax 0.9.2
jaxlib 0.9.2
jax-cuda12-plugin 0.9.2
jax-cuda12-pjrt 0.9.2
**Environment**:
```
export CUDA_ROOT=$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia
export LD_LIBRARY_PATH=$CUDA_ROOT/cuda_runtime/lib
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CUDA_ROOT/cublas/lib
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CUDA_ROOT/cudnn/lib
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CUDA_ROOT/cuda_nvrtc/lib
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu
```
**Reproduce**:
```
python3 -c "import jax; print(jax.devices()); print(jax.default_backend())"
```
**Error**:
FAILED_PRECONDITION: No visible GPU devices
**Verification**:
- nvidia-smi works
- ctypes.CDLL("libcuda.so").cuInit(0) returns 0
- libcuda, libcudart, libcudnn, libnvidia-ml all load correctly
**Question**:
- Is RTX 50-series (Blackwell) supported in jax 0.9.x CUDA builds?
- If not, what is the recommended setup for Python 3.12 + CUDA + Blackwell?
Contributor guide
Assessment
This issue has not been assessed yet.