{"record":{"id":"71b84b52b9068de8","repo":"jax-ml/jax","slug":"sharding-on-sequence-dim-is-not-allowed","errorCode":null,"errorMessage":"Sharding on sequence dim is not allowed.","messagePattern":"Sharding on sequence dim is not allowed\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/fused_attention_stablehlo.py","lineNumber":988,"sourceCode":"def _get_padded_spec(arg_info):\n  spec = None if arg_info.sharding is None else arg_info.sharding.spec\n  ndim = arg_info.ndim\n  if spec is None:\n    return (None,) * ndim\n  assert len(spec) <= ndim\n  return spec + (None,) * (ndim - len(spec))\n\ndef _check_qkv_bias_mask_spec(\n    query_spec, key_spec, value_spec, bias_spec, layout):\n  # check qkv spec\n  if not query_spec == key_spec == value_spec:\n    raise ValueError(\"Query, key and value should have same sharding.\")\n  if layout == AttentionLayout.BNTH.value:\n    *batch_spec, num_head_spec, q_seq_spec, head_spec = query_spec\n  else:\n    *batch_spec, q_seq_spec, num_head_spec, head_spec = query_spec\n  if q_seq_spec is not None:\n    raise ValueError(\"Sharding on sequence dim is not allowed.\")\n  if head_spec is not None:\n    raise ValueError(\"Sharding on head dim is not allowed.\")\n  # check bias spec\n  if bias_spec:\n    *bias_batch_spec, bias_num_head_spec, bias_q_seq_spec, bias_kv_seq_spec = bias_spec\n    if any(bias_batch_spec) and bias_batch_spec != batch_spec or \\\n      bias_num_head_spec is not None and bias_num_head_spec != num_head_spec:\n      raise ValueError(\n        \"Query and bias should have same sharding on batch and num_head dim.\")\n    if bias_q_seq_spec is not None or bias_kv_seq_spec is not None:\n      raise ValueError(\"Sharding on bias sequence dim is not allowed.\")\n\n\n# fwd custom partition\ndef _infer_fwd_output_sharding(mesh, arg_shapes, variadic_args, is_training, layout):\n  # only sharding on batch and num_head dim is allowed\n  # (*batch, q_seq, num_head, head)\n  query_spec = _get_padded_spec(arg_shapes[0])","sourceCodeStart":970,"sourceCodeEnd":1006,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/fused_attention_stablehlo.py#L970-L1006","documentation":"Raised by _check_qkv_bias_mask_spec when the query's sharding spec places a mesh axis on the sequence dimension (q_seq) for either BNTH or BHTH layout. The cuDNN fused attention SPMD support only allows sharding on batch and num_heads dimensions, because the kernel treats sequence as a local dimension.","triggerScenarios":"Passing a PartitionSpec with a non-None entry in the sequence position, e.g. P('batch', 'seq', 'heads', None) for (batch, seq, heads, head_dim), under jitted sharded dot_product_attention.","commonSituations":"Sequence-parallel model code ported to the fused attention API; sharding long-context activations over sequence to fit memory; assuming FSDP-style full sharding works on every axis.","solutions":["Remove the mesh axis from the sequence dim: shard on batch and/or num_heads only, e.g. P('batch', None, 'heads', None)","For long sequences use sharded-batch data parallelism or sequence chunking outside the fused call","Wrap with jax.lax.with_sharding_constraint to reshard to an allowed spec before the attention call"],"exampleFix":"# before\nq_spec = P('batch', 'seq', 'heads', None)  # sequence sharded -> error\n\n# after\nq_spec = P('batch', None, 'heads', None)  # shard batch & heads only","handlingStrategy":"validation","validationCode":"from jax.sharding import PartitionSpec\ndef spec_ok_for_attention(spec, layout='BNTH'):\n    # BNTH: (batch..., seq, heads, head); BHTH: (batch..., heads, seq, head)\n    seq_idx = -3 if layout == 'BNTH' else -2\n    return spec[seq_idx] is None and spec[-1] is None","typeGuard":"def attention_spec_valid(spec: PartitionSpec, layout: str) -> bool:\n    seq_idx = -3 if layout == 'BNTH' else -2\n    return spec[seq_idx] is None and spec[-1] is None","tryCatchPattern":null,"preventionTips":["Restrict attention operand shardings to batch and num_heads mesh axes","Add unit tests asserting specs of the form P('batch', None, 'heads', None)"],"tags":["jax","sharding","sequence-parallel","spmd","cudnn"],"backgroundTag":"sharding-disallowed-axis","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}