GPU detected successfully but first CUDA allocation fails with CUDA_ERROR_UNKNOWN on WSL2 (RTX A4000)

Description

Summary

JAX successfully detects the NVIDIA GPU and selects the GPU backend.

import jax

print(jax.devices())
print(jax.default_backend())

Output:

[CudaDevice(id=0)]
gpu

However, the first GPU memory allocation fails.

Minimal reproducer:

import jax.numpy as jnp

print(jnp.array([1, 2, 3]))

Result:

Failed to allocate device memory of 2.2KiB (2304 bytes)

INTERNAL:
CUDA error:
CUDA_ERROR_UNKNOWN

The failure occurs on the first allocation attempt.

---

Backend Verification

import jax

print("backend =", jax.default_backend())
print("devices =", jax.devices())
print("gpu devices =", jax.devices("gpu"))

Output:

backend = gpu
devices = [CudaDevice(id=0)]
gpu devices = [CudaDevice(id=0)]

This confirms that JAX successfully detects the GPU and selects the GPU backend before the allocation failure occurs.

---

Reproducer

import jax
import jax.numpy as jnp

print("jax =", jax.__version__)
print("devices =", jax.devices())
print("backend =", jax.default_backend())

x = jnp.array([1, 2, 3])

print(x)

---

Observed Behavior

Output:

jax = 0.10.2
devices = [CudaDevice(id=0)]
backend = gpu

E.... cuda_executor.cc:1182
Failed to allocate device memory of 2.2KiB (2304 bytes)

INTERNAL:
CUDA error:
CUDA_ERROR_UNKNOWN

The error is repeatedly emitted from:

cuda_executor.cc:1182

---

Additional Test

Explicit device placement also fails.

import jax
import jax.numpy as jnp

gpu = jax.devices("gpu")[0]

x = jax.device_put(
    jnp.array([1, 2, 3]),
    device=gpu,
)

print(x)

Output:

E.... cuda_executor.cc:1182
Failed to allocate device memory of 2.2KiB (2304 bytes)

INTERNAL:
CUDA error:
CUDA_ERROR_UNKNOWN

This indicates that:

• JAX can discover the GPU device.
• jax.devices("gpu") succeeds.
• Explicit GPU device selection succeeds.
• The failure occurs when JAX attempts the first GPU memory allocation.

---

Control Experiments

PyTorch CUDA Works

import torch

print(torch.cuda.is_available())

x = torch.tensor([1, 2, 3], device="cuda")

print(x)

Output:

True
tensor([1, 2, 3], device='cuda:0')

This suggests that:

• GPU access works
• CUDA driver works
• WSL GPU virtualization works
• GPU memory allocation works through PyTorch

CPU Backend Works

export JAX_PLATFORMS=cpu

import jax
import jax.numpy as jnp

print(jax.devices())
print(jnp.array([1, 2, 3]))

Output:

[CpuDevice(id=0)]
[1 2 3]

---

CUDA Library Checks

ls -l /usr/lib/wsl/lib/libcuda*

Output:

/usr/lib/wsl/lib/libcuda.so
/usr/lib/wsl/lib/libcuda.so.1
/usr/lib/wsl/lib/libcuda.so.1.1

Library load test:

import ctypes

ctypes.CDLL("libcuda.so")
print("success")

Output:

success

---

Additional Observation

Running with:

JAX_PLATFORMS=gpu

results in:

RuntimeError:
Unable to initialize backend 'rocm':
Backend 'rocm' is not in the list of known backends:
['cpu', 'tpu', 'cuda']

even though a normal JAX session reports:

[CudaDevice(id=0)]

and

backend = gpu

I am not sure whether this observation is related to the allocation failure, but it may be relevant.

---

Conclusion

• JAX successfully detects the GPU.
• jax.devices() returns [CudaDevice(id=0)].
• jax.devices("gpu") returns [CudaDevice(id=0)].
• The default backend is gpu.
• The first GPU memory allocation fails with CUDA_ERROR_UNKNOWN.
• The failure occurs even for a very small allocation (2304 bytes).
• Explicit jax.device_put(..., device=gpu) fails with the same error.
• PyTorch CUDA works correctly on the same machine.
• CPU JAX works correctly.
• CUDA libraries appear to be visible from WSL.

Since PyTorch can allocate CUDA memory successfully on the same system, this does not appear to be a general CUDA driver, GPU hardware, or WSL GPU virtualization issue.

The problem appears to be specific to the JAX GPU backend initialization and/or allocation path under this WSL2 environment.

Any suggestions for additional diagnostics would be appreciated.

System info (python version, jaxlib version, accelerator, etc.)

OS: Windows 11 host WSL2 Ubuntu 24.04.4 LTS (Noble Numbat)

Conda Environment: colabfold_gpu

Python: 3.11.16

JAX: jax==0.10.2 jaxlib==0.10.2 jax-cuda12-plugin==0.10.2 jax-cuda12-pjrt==0.10.2

Backend: gpu

Devices: [CudaDevice(id=0)] jax.devices("gpu") -> [CudaDevice(id=0)]

GPU: NVIDIA RTX A4000 16 GB VRAM

NVIDIA Driver: Driver Version: 596.71 CUDA Version: 13.2

NVIDIA-SMI: NVIDIA-SMI 595.83.02

GPU Visibility: nvidia-smi -L -> GPU 0: NVIDIA RTX A4000

JAX installation:

conda list | grep jax

jax 0.10.2 pypi_0 jax-cuda12-pjrt 0.10.2 pypi_0 jax-cuda12-plugin 0.10.2 pypi_0 jaxlib 0.10.2 pypi_0

Relevant dependencies:

numpy==2.4.6 scipy==1.17.1 ml_dtypes==0.6.0

nvidia-cublas-cu12==12.9.2.10 nvidia-cuda-runtime-cu12==12.9.79 nvidia-cudnn-cu12==9.25.1.1 nvidia-cufft-cu12==11.4.1.4 nvidia-cusolver-cu12==11.7.5.82 nvidia-cusparse-cu12==12.5.10.65 nvidia-nccl-cu12==2.31.2

Observed Failure:

Failed to allocate device memory of 2.2KiB (2304 bytes)

INTERNAL: CUDA error: CUDA_ERROR_UNKNOWN