jax-ml/jax · error · NotImplementedError

Packed layout requires a GPU with at least Hopper architectu

Error message

Packed layout requires a GPU with at least Hopper architecture.

What it means

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.

Source

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

        # Check the head dim.
        is_hopper_or_later = check_compute_capability("9.0")
        H_max = 256 if is_hopper_or_later else 128
        # 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

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

# before
out = jax.nn.dot_product_attention(q, k, v, q_seqlen=q_offsets, kv_seqlen=kv_offsets)

# after
cap = float(jax.devices()[0].compute_capability)
if cap >= 9.0:
  out = jax.nn.dot_product_attention(q, k, v, q_seqlen=q_offsets, kv_seqlen=kv_offsets)
else:
  out = padded_or_per_sequence_attention(q, k, v, lengths)  # fallback path

Defensive patterns

Strategy: type-guard

Validate before calling

import jax
def supports_packed():
    return tuple(jax.devices()[0].compute_capability) >= (9, 0)

Type guard

def is_hopper_or_newer(dev) -> bool:
    cc = tuple(int(x) for x in dev.compute_capability)
    return cc >= (9, 0)

Prevention

Gate packed-layout attention behind a compute-capability check at model-construction time
Document GPU requirements (Hopper+) in configs that enable packed sequences

When it happens

Trigger: 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.

Common situations: 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.

Related errors

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