google-deepmind / google-deepmind/mujoco_playground

JAX 0.9.x CUDA: "No visible GPU devices" on RTX 5090 (Blackwell)

Open
#299 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.