{"record":{"id":"0344ded4061746f0","repo":"jax-ml/jax","slug":"mla-requires-cudnn-version-9-10-and-at-least-ho","errorCode":null,"errorMessage":"mla requires cudnn version >= 9.10 and at least hopper arch.","messagePattern":"mla requires cudnn version >= 9\\.10 and at least hopper arch\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/fused_attention_stablehlo.py","lineNumber":401,"sourceCode":"        # check if multi-head latent attention is needed\n        is_mla = qH != vH\n        if not (qH <= H_max and qH % 8 == 0):\n          raise NotImplementedError(\n              f\"The head dim must be <= {H_max} and a multiple of 8, \"\n              f\"but got {qH}.\"\n          )\n\n        # Check patterns with bias, seqlen should be divisible by 2\n        if (is_training and has_bias and (T % 2 != 0 or S % 2 != 0)):\n          raise NotImplementedError(\n              f\"Unsupported sequence length Q {T}, KV {S}.\"\n          )\n\n        if is_packed and  not check_compute_capability(\"9.0\"):\n          raise NotImplementedError(\n            \"Packed layout requires a GPU with at least Hopper architecture.\")\n        if is_mla and (cudnn_version < 91000 or not check_compute_capability(\"9.0\")):\n          raise NotImplementedError(\n            \"mla requires cudnn version >= 9.10 and at least hopper arch.\")\n\ndef check_cudnn_version():\n  # check if cuDNN is installed\n  if cuda_versions is None:\n    raise RuntimeError(\"cuDNN is not detected.\")\n  return cuda_versions.cudnn_get_version()\n\ndef check_compute_capability(capability):\n  if not 'cuda' in xla_bridge.get_backend().platform_version:\n    return False\n  d, *_ = xla_bridge.local_devices(backend=\"gpu\")\n  target = tuple(int(x) for x in capability.split(\".\"))\n  current = tuple(int(x) for x in d.compute_capability.split(\".\"))\n  return current >= target\n\ndef is_cuda_compute_capability_equal(capability):\n  if not 'cuda' in xla_bridge.get_backend().platform_version:","sourceCodeStart":383,"sourceCodeEnd":419,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/fused_attention_stablehlo.py#L383-L419","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","solutions":["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","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)","Verify versions first: check cudnn version via jax's cuda_versions and device compute capability before choosing MLA layout"],"exampleFix":"# before\nout = jax.nn.dot_product_attention(q, k, v, ...,)  # mla layout selected -> error\n\n# after\nfrom jax._src.cudnn.fused_attention_stablehlo import check_cudnn_version, check_compute_capability\nuse_mla = check_cudnn_version() >= 91000 and check_compute_capability(\"9.0\")\nout = mla_fused_attention(q, k, v) if use_mla else reference_mla(q, k, v)","handlingStrategy":"type-guard","validationCode":"from jax._src.cudnn.fused_attention_stablehlo import check_cudnn_version, check_compute_capability\ndef mla_supported():\n    return check_cudnn_version() >= 91000 and check_compute_capability(\"9.0\")","typeGuard":"def can_use_mla() -> bool:\n    try:\n        return check_cudnn_version() >= 91000 and check_compute_capability(\"9.0\")\n    except RuntimeError:\n        return False","tryCatchPattern":null,"preventionTips":["Pin cuDNN >= 9.10 in environment files when MLA is used","Add a startup assertion logging cuDNN version and GPU compute capability"],"tags":["jax","cudnn","mla","version-mismatch","gpu-architecture"],"backgroundTag":"library-version-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}