Per-operator (non-jitted) kernels return wrong values / NaNs on RTX A6000 (CUDA 13), jitted path and CPU are correct; breaks a jitted training loop

Description

On two NVIDIA RTX A6000 instances (Ubuntu 24.04, kernel 6.17.0-35, driver 610.43.02, CUDA UMD 13.3), with jax / jaxlib 0.11.1, several per-operator (non-jitted) kernels return wrong values or NaNs, while the same expressions under jax.jit produce correct results on the same device, and the same per-operator code on CPU is entirely correct. The same misbehavior reproduces identically across two rented instances and across both the jax[cuda13] and jax[cuda12] wheels. It also breaks a small training loop: the first jitted training step produces NaNs on GPU while the identical loop on CPU trains normally.

Per-operator vs jit (probe3_minimal.py, attached)

Per-operator (under jax.disable_jit()), run on two instances (ddznt3yn, ae0rs8dp), field-identical outputs:

op	result
jnp.sqrt(512.)	0.0 (expected 22.627)
jnp.ones(4) + 1	[513, 513, 513, 513] (expected 2.0)
jnp.arange(10).sum()	45 (correct)
ones(4,4) @ ones(4,4)	4.0 (correct)
jax.random.normal(key, (4,)) … (16384000,)	all-NaN at every size
jax.random.uniform(key, (1000,))	6 NaNs; head [2.53e+30, 1.42e+35, 1.4142135, 1.4142135]
jax.random.randint(key, (1000,), 0, 32000)	no NaN but wrong values, e.g. -645928121, -1274445908, -822036043, -74888185

Under jax.jit, the same five ops are all correct, and CPU per-operator is correct across all 16 rows. Two adjacent statements suggest reads crossing statement boundaries: sqrt(512) loses its input and returns 0.0, and the following ones(4)+1 returns 513 = 1 + 512 — the previous statement's operand. We only report the reading; we are not in a position to diagnose the mechanism.

The jax[cuda12] wheel reproduces the same table, except uniform(1000) returns 0 NaNs with head values all 2.8284270763397217.

A separate, earlier process on the same instance (probe_opbyop_ddznt3yn.log, two reps) shows a different bad-value pattern: normal((32000, 512)) all zeros (max_abs 0.0), then s/sqrt(d) produces 262144 NaNs, while jitted softmax is fine (max_abs 0.0078125). So the bad values differ across processes.

Training-loop break (minimal file, init outside jit, only the step is jitted)

• GPU: loss [10.373490333557129, NaN, NaN]
• CPU: [10.373490333557129, 10.37348461151123, 10.37350845336914, 10.373493194580078] (two identical CPU runs)
• With JAX_DEBUG_NANS=1: stdout prints twice "Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version."; stderr raises FloatingPointError: invalid value (nan) encountered in div at ufuncs.py:2478 (lax.div), pointing at the softmax line att = jax.nn.softmax((q @ k.transpose(0, 2, 1)) / jnp.sqrt(DMODEL), axis=-1).
• The jax[cuda12] wheel reproduces the same [10.37349, NaN, NaN].
• Putting initialization inside the jit too does not recover: three runs all [10.373491287231445, NaN, NaN] (with --xla_gpu_exclude_nondeterministic_ops=true).

Reproduction scope

• Training-loop NaN reproduced on four rented A6000 instances (2026-09-09 through 2026-09-11 UTC), same image; the per-operator probe reproduced on two of them. Instances have been released, so we cannot rerun on request, but every log and artifact is attached.

Search

We searched before filing; queries actually used: nan disable_jit, disable_jit NaN gpu, op-by-op NaN, eager execution NaN random, jnp.random.normal NaN gpu, disable_jit zeros, eager NaN, un-jitted NaN, op-by-op, eager mode, sqrt 512 NaN, A6000 jax NaN. The closest match we found is #25701 (closed 2025-04), which describes a similar-looking "nan encountered in jit even when jit disabled" error message, but the cause and context there are different (a log of a negative constant on CPU with debug_nans); ours is wrong values and NaNs from per-operator GPU kernels, without debug flags.

System info

env_ae0rs8dp.txt (full output of nvidia-smi + jax env)

jax:    0.11.1
jaxlib: 0.11.1
numpy:  2.5.3
python: 3.12.13
device info: NVIDIA RTX A6000-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='instance-ae0rs8dp-main', release='6.17.0-35-generic', version='#35~24.04.1-Ubuntu SMP PREEMPT_DYNAMIC Tue May 26 19:30:42 UTC 2', machine='x86_64')

$ nvidia-smi
NVIDIA-SMI 610.43.02              KMD Version: 610.43.02     CUDA UMD Version: 13.3
GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
   0  NVIDIA RTX A6000               Off |   00000000:04:00.0 Off |                  Off |
Driver Version         : 610.43.02 [Deprecated; will be removed in CUDA 14.0. Use KMD Version instead]
CUDA Version           : 13.3 [Deprecated; will be removed in CUDA 14.0. Use CUDA UMD Version instead]

(Full env_ae0rs8dp.txt including the two nvidia-smi "Deprecated" annotations is attached in the gist.)

Gist with all scripts and logs: https://gist.github.com/KitKyoD/3e1c04633e476aec53078e2f660a3df7