jax-ml/jax · error · NotImplementedError

mla requires cudnn version >= 9.10 and at least hopper arch.

Error message

mla requires cudnn version >= 9.10 and at least hopper arch.

What it means

Raised by check_is_flash_attention when the MLA (multi-head latent attention) layout is requested but cuDNN is older than 9.10 or the GPU is pre-Hopper. cuDNN's MLA fused attention kernels require cudnn >= 9.10 and compute capability >= 9.0 (Hopper).

Source

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

        # check if multi-head latent attention is needed
        is_mla = qH != vH
        if not (qH <= H_max and qH % 8 == 0):
          raise NotImplementedError(
              f"The head dim must be <= {H_max} and a multiple of 8, "
              f"but got {qH}."
          )

        # 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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Upgrade cuDNN to >= 9.10 (e.g. pip install -U nvidia-cudnn-cu12 or install a JAX CUDA wheel bundling cuDNN 9.10+) and run on a Hopper GPU
2. If the GPU is pre-Hopper, use the non-fused attention implementation for MLA (fall back to standard dot_product_attention math or the Flax attention path)
3. Verify versions first: check cudnn version via jax's cuda_versions and device compute capability before choosing MLA layout

Example fix

# before
out = jax.nn.dot_product_attention(q, k, v, ...,)  # mla layout selected -> error

# after
from jax._src.cudnn.fused_attention_stablehlo import check_cudnn_version, check_compute_capability
use_mla = check_cudnn_version() >= 91000 and check_compute_capability("9.0")
out = mla_fused_attention(q, k, v) if use_mla else reference_mla(q, k, v)

from jax._src.cudnn.fused_attention_stablehlo import check_cudnn_version, check_compute_capability
def mla_supported():
    return check_cudnn_version() >= 91000 and check_compute_capability("9.0")

def can_use_mla() -> bool:
    try:
        return check_cudnn_version() >= 91000 and check_compute_capability("9.0")
    except RuntimeError:
        return False

When it happens

Trigger: Calling jax.nn.dot_product_attention with the MLA layout on a system with cuDNN < 9.10 (e.g. 9.6 shipped with CUDA 12.5) or on any pre-Hopper GPU.

Common situations: Running DeepSeek-V3/R1-style MLA models on older CUDA/cuDNN stacks or A100 nodes; docker images with stale cuDNN; upgrading JAX without upgrading the CUDA/cuDNN wheels.

Related errors

• Packed layout requires a GPU with at least Hopper architectu
• stop_gradient only works on valid JAX arrays, but input argu
• Unknown GPU platform for __dlpack__: {platform_version}
• Unsupported sequence length Q {T}, KV {S}.
• cuDNN is not detected.

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