{"record":{"id":"45f17d179cb4a6b2","repo":"jax-ml/jax","slug":"unsupported-sequence-length-q-t-kv-s","errorCode":null,"errorMessage":"Unsupported sequence length Q {T}, KV {S}.","messagePattern":"Unsupported sequence length Q (.+?), KV (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/fused_attention_stablehlo.py","lineNumber":393,"sourceCode":"            raise NotImplementedError(\n                f\"Unsupported sequence length Q {T}, KV {S} and head dim {qH} for FP8.\"\n            )\n    else:\n        # bf16/fp16 attention conditions\n        # 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:","sourceCodeStart":375,"sourceCodeEnd":411,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/fused_attention_stablehlo.py#L375-L411","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Pad Q and KV sequence lengths to the next even number (e.g. pad to multiple of 2) before calling dot_product_attention","Use an odd-length-safe path: pass bias=None, or disable cuDNN fusion so JAX falls back to the standard attention implementation","If odd lengths are essential during training, run without bias and add bias via a separate masked softmax step outside the fused kernel"],"exampleFix":"# before\nattn = jax.nn.dot_product_attention(q, k, v, bias=bias, is_training=True)  # T=127 -> error\n\n# after\npad = T % 2\nq = jax.lax.pad(q, 0.0, [(0,0,0),(0,pad,0),(0,0,0),(0,0,0)])\nk = jax.lax.pad(k, 0.0, [(0,0,0),(0,pad,0),(0,0,0),(0,0,0)])\nv = jax.lax.pad(v, 0.0, [(0,0,0),(0,pad,0),(0,0,0),(0,0,0)])\nattn = jax.nn.dot_product_attention(q, k, v, bias=bias, is_training=True)[:, :T]","handlingStrategy":"validation","validationCode":"def check_bias_seqlens(T, S, has_bias, is_training):\n    if is_training and has_bias and (T % 2 or S % 2):\n        raise ValueError(f\"Pad Q ({T}) and KV ({S}) to even lengths for cuDNN fused attention with bias.\")","typeGuard":null,"tryCatchPattern":"try:\n    out = jax.nn.dot_product_attention(q, k, v, bias=bias, is_training=True)\nexcept NotImplementedError:\n    out = manual_attention(q, k, v, bias)","preventionTips":["Always round sequence lengths up to an even number when bias is used in training","Assert q.shape[-3] % 2 == 0 and kv.shape[-3] % 2 == 0 in data pipelines feeding fused attention"],"tags":["jax","cudnn","flash-attention","sequence-length","training"],"backgroundTag":"unsupported-shape-constraint","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}