jax-ml/jax · critical · RuntimeError

cuDNN is not detected.

Error message

cuDNN is not detected.

What it means

Raised by check_cudnn_version when JAX cannot detect a cuDNN installation at all (cuda_versions is None). The cuDNN fused attention feature requires cuDNN; without it the API refuses to run instead of failing deep inside XLA.

Source

Thrown at jax/_src/cudnn/fused_attention_stablehlo.py:407

          )

        # Check patterns with bias, seqlen should be divisible by 2
        if (is_training and has_bias and (T % 2 != 0 or S % 2 != 0)):
          raise NotImplementedError(
              f"Unsupported sequence length Q {T}, KV {S}."
          )

        if is_packed and  not check_compute_capability("9.0"):
          raise NotImplementedError(
            "Packed layout requires a GPU with at least Hopper architecture.")
        if is_mla and (cudnn_version < 91000 or not check_compute_capability("9.0")):
          raise NotImplementedError(
            "mla requires cudnn version >= 9.10 and at least hopper arch.")

def check_cudnn_version():
  # check if cuDNN is installed
  if cuda_versions is None:
    raise RuntimeError("cuDNN is not detected.")
  return cuda_versions.cudnn_get_version()

def check_compute_capability(capability):
  if not 'cuda' in xla_bridge.get_backend().platform_version:
    return False
  d, *_ = xla_bridge.local_devices(backend="gpu")
  target = tuple(int(x) for x in capability.split("."))
  current = tuple(int(x) for x in d.compute_capability.split("."))
  return current >= target

def is_cuda_compute_capability_equal(capability):
  if not 'cuda' in xla_bridge.get_backend().platform_version:
    return False
  d, *_ = xla_bridge.local_devices(backend="gpu")
  target = tuple(int(x) for x in capability.split("."))
  current = tuple(int(x) for x in d.compute_capability.split("."))
  return current == target

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Install GPU-enabled JAX with bundled deps: pip install -U 'jax[cuda12]'
2. If using pip CUDA wheels, add nvidia-cudnn-cu12 and ensure its lib dir is on LD_LIBRARY_PATH (or rely on the jax[cuda12] extra)
3. Verify detection: python -c 'from jax._src.cudnn.fused_attention_stablehlo import check_cudnn_version; print(check_cudnn_version())'
4. If running a JAX build that supports a different CUDA version, align CUDA/cuDNN/JAX versions per the JAX installation guide

Example fix

# before: jax installed CPU-only -> RuntimeError: cuDNN is not detected.
pip install jax

# after
pip install -U "jax[cuda12]"

from jax._src.cudnn.fused_attention_stablehlo import check_cudnn_version
def require_cudnn():
    try:
        return check_cudnn_version()
    except RuntimeError as e:
        raise RuntimeError(f"Install jax[cuda12] with cuDNN: {e}") from e

try:
    check_cudnn_version()
except RuntimeError:
    # fall back to non-cuDNN attention path
    out = fallback_attention(q, k, v)

When it happens

Trigger: Calling jax.nn.dot_product_attention (cuDNN path), paged_attention, or the related tests on a machine where cuDNN is not installed or not discoverable by JAX (CPU-only JAX, missing nvidia-cudnn wheel, broken LD_LIBRARY_PATH).

Common situations: Installed jax (CPU) instead of jax[cuda12]/jax_cuda12_plugin; pip CUDA package set without the matching nvidia-cudnn-cu12 wheel; conda environments with cuDNN not on the loader path; slim docker images.

Related errors

• Couldn't get local_hardware_id for __dlpack__
• jax.pure_callback failed to find a local CPU device to place
• Unsupported sequence length Q {T}, KV {S}.
• Packed layout requires a GPU with at least Hopper architectu
• mla requires cudnn version >= 9.10 and at least hopper arch.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/9501ddedbde008bc. Report an issue: GitHub.