jax-ml/jax · error · NotImplementedError

Unsupported sequence length Q {T}, KV {S}.

Error message

Unsupported sequence length Q {T}, KV {S}.

What it means

Raised by check_is_flash_attention in JAX's cuDNN fused attention StableHLO path when training with an attention bias and either the query or KV sequence length is odd. cuDNN's fused flash attention kernels with bias require sequence lengths divisible by 2 during training, so JAX refuses the configuration rather than producing wrong results.

Source

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

            raise NotImplementedError(
                f"Unsupported sequence length Q {T}, KV {S} and head dim {qH} for FP8."
            )
    else:
        # bf16/fp16 attention conditions
        # 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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pad Q and KV sequence lengths to the next even number (e.g. pad to multiple of 2) before calling dot_product_attention
  2. Use an odd-length-safe path: pass bias=None, or disable cuDNN fusion so JAX falls back to the standard attention implementation
  3. If odd lengths are essential during training, run without bias and add bias via a separate masked softmax step outside the fused kernel

Example fix

# before
attn = jax.nn.dot_product_attention(q, k, v, bias=bias, is_training=True)  # T=127 -> error

# after
pad = T % 2
q = jax.lax.pad(q, 0.0, [(0,0,0),(0,pad,0),(0,0,0),(0,0,0)])
k = jax.lax.pad(k, 0.0, [(0,0,0),(0,pad,0),(0,0,0),(0,0,0)])
v = jax.lax.pad(v, 0.0, [(0,0,0),(0,pad,0),(0,0,0),(0,0,0)])
attn = jax.nn.dot_product_attention(q, k, v, bias=bias, is_training=True)[:, :T]
Defensive patterns

Strategy: validation

Validate before calling

def check_bias_seqlens(T, S, has_bias, is_training):
    if is_training and has_bias and (T % 2 or S % 2):
        raise ValueError(f"Pad Q ({T}) and KV ({S}) to even lengths for cuDNN fused attention with bias.")

Try / catch

try:
    out = jax.nn.dot_product_attention(q, k, v, bias=bias, is_training=True)
except NotImplementedError:
    out = manual_attention(q, k, v, bias)

Prevention

When it happens

Trigger: Calling jax.nn.dot_product_attention (or the cuDNN fused attention path) with bias/mask not None, is_training=True, and q_seq_len % 2 != 0 or kv_seq_len % 2 != 0.

Common situations: Training a model with additive attention bias (e.g. T5-style relative position bias, ALiBi) on odd sequence lengths like 127, 511, or unpadded variable-length batches.

Related errors

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