{"record":{"id":"504eda759ef979af","repo":"jax-ml/jax","slug":"packed-layout-requires-a-gpu-with-at-least-hopper","errorCode":null,"errorMessage":"Packed layout requires a GPU with at least Hopper architecture.","messagePattern":"Packed layout requires a GPU with at least Hopper architecture\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/fused_attention_stablehlo.py","lineNumber":398,"sourceCode":"        # Check the head dim.\n        is_hopper_or_later = check_compute_capability(\"9.0\")\n        H_max = 256 if is_hopper_or_later else 128\n        # 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","sourceCodeStart":380,"sourceCodeEnd":416,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/fused_attention_stablehlo.py#L380-L416","documentation":"Raised by check_is_flash_attention when the packed layout (nvdim=2 packed sequences with seq_offsets) is requested but the GPU is not at least NVIDIA Hopper (compute capability 9.0). cuDNN's packed/varlen attention layout relies on kernels only available on Hopper and newer architectures.","triggerScenarios":"Calling jax.nn.dot_product_attention with packed sequence offsets (q_seqlen/kv_seqlen argument, layout with 2D batch/seq dims) on an Ampere (A100, compute 8.0) or older GPU.","commonSituations":"Developing packed-attention code on A100 or consumer Ampere GPUs, or running a checkpoint/config tuned for H100 on older cluster nodes; CI machines with older GPUs.","solutions":["Run on a Hopper (H100/H200) or newer GPU (compute capability >= 9.0)","If stuck on pre-Hopper hardware, unpad/de-pad sequences and call attention per-sequence or with explicit padding masks instead of packed layout","Gate the packed path at config level: check jax devices' compute capability before selecting layout"],"exampleFix":"# before\nout = jax.nn.dot_product_attention(q, k, v, q_seqlen=q_offsets, kv_seqlen=kv_offsets)\n\n# after\ncap = float(jax.devices()[0].compute_capability)\nif cap >= 9.0:\n  out = jax.nn.dot_product_attention(q, k, v, q_seqlen=q_offsets, kv_seqlen=kv_offsets)\nelse:\n  out = padded_or_per_sequence_attention(q, k, v, lengths)  # fallback path","handlingStrategy":"type-guard","validationCode":"import jax\ndef supports_packed():\n    return tuple(jax.devices()[0].compute_capability) >= (9, 0)","typeGuard":"def is_hopper_or_newer(dev) -> bool:\n    cc = tuple(int(x) for x in dev.compute_capability)\n    return cc >= (9, 0)","tryCatchPattern":null,"preventionTips":["Gate packed-layout attention behind a compute-capability check at model-construction time","Document GPU requirements (Hopper+) in configs that enable packed sequences"],"tags":["jax","cudnn","gpu-architecture","packed-attention","hopper"],"backgroundTag":"gpu-capability-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}