[cuda] Bump nvidia-cuda-nvcc-cu12 dependency to 12.6.85

[gspschmid]

Bumps the minimum version for nvidia-cuda-nvcc-cu12 to 12.6.85, which is the earliest published version that of that wheel incorporating CUDA 12.6.3.

This resolves the issue underlying #24438. Here's a simplified reproducer:

$ docker run -ti --gpus all ubuntu
$ nvidia-smi
(... H100)

$ apt-get update
$ apt-get install -y python3-pip
$ export PIP_BREAK_SYSTEM_PACKAGES=1

# Install JAX nightly (picks up recent nvidia-cuda-nvcc-cu12 version, e.g. 12.6.85)
$ pip install -U --pre jax jaxlib jax-cuda12-plugin[with_cuda] jax-cuda12-pjrt -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html

# => GOOD
$ python3 -c "import jax; import jax.numpy as jnp; A = jnp.arange(18).reshape(6, 3); m = jnp.arange(-3, 3); print(jax.jit(lambda _0, _1: _0.at[jnp.abs(_1), 0].get())(A, m))"
[9 6 3 0 3 6]

# Force-install bad version
$ pip install --force-reinstall "nvidia-cuda-nvcc-cu12==12.6.77"

# => BAD
$ python3 -c "import jax; import jax.numpy as jnp; A = jnp.arange(18).reshape(6, 3); m = jnp.arange(-3, 3); print(jax.jit(lambda _0, _1: _0.at[jnp.abs(_1), 0].get())(A, m))"
[0 0 0 0 3 6]

# Upgrade nvcc dependency
$ pip install -U "nvidia-cuda-nvcc-cu12>=12.6.85"

# => GOOD
$ python3 -c "import jax; import jax.numpy as jnp; A = jnp.arange(18).reshape(6, 3); m = jnp.arange(-3, 3); print(jax.jit(lambda _0, _1: _0.at[jnp.abs(_1), 0].get())(A, m))"
[9 6 3 0 3 6]

As a follow-up we might also want to add a warning asking users with an affected version of nvidia-cuda-nvcc-cu12 to upgrade (e.g., via pip install -U "nvidia-cuda-nvcc-cu12>=12.6.85").

@hawkinsp Any opinions on whether to have that warning / where it should go? As to how to check for the current ptxas version (https://github.com/jax-ml/jax/pull/24438/files#r1818345222), I suppose we might expose that alongside other CUDA versions here (

jax/jaxlib/cuda/versions.cc

Line 26 in c35f8b2

NB_MODULE(_versions, m) {

)? A cruder alternative approach might be to check for the nvidia-cuda-nvcc-cu12 wheel's __version__, if present.

[hawkinsp]

I think the best place for the warning would be in XLA: it should warn if it detects a ptxas that is known to be buggy.

[gspschmid]

Corresponding XLA PR that emits a warning: openxla/xla#19927