{"record":{"id":"f2c112a988c1c698","repo":"jax-ml/jax","slug":"sharding-on-head-dim-is-not-allowed","errorCode":null,"errorMessage":"Sharding on head dim is not allowed.","messagePattern":"Sharding on head dim is not allowed\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/fused_attention_stablehlo.py","lineNumber":990,"sourceCode":"  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])\n  # (*batch, kv_seq, num_head, head)\n  key_spec = _get_padded_spec(arg_shapes[1])","sourceCodeStart":972,"sourceCodeEnd":1008,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/fused_attention_stablehlo.py#L972-L1008","documentation":"Raised by _check_qkv_bias_mask_spec when the query's sharding spec places a mesh axis on the per-head feature (head) dimension. cuDNN fused attention's custom partitioner only supports sharding on batch and num_heads; the head dimension must stay replicated on each device.","triggerScenarios":"Passing a PartitionSpec with a mesh name in the last (head_dim) position, e.g. P('batch', None, 'heads', 'model') under jitted sharded dot_product_attention.","commonSituations":"Reusing tensor-parallel weight shardings (which shard the head dim of WQ/WK/WV outputs) directly on the attention operands; applying Megatron-style TP specs unchanged to this API.","solutions":["Use a sharding with None in the last position, e.g. P('batch', None, 'heads', None) or P(None, None, 'tp', None)","If tensor-parallel over head_dim is required, implement attention manually (matmul/softmax) instead of the cuDNN fused path","Reshard q/k/v to an allowed spec before the call via with_sharding_constraint"],"exampleFix":"# before\nq_spec = P('batch', None, 'heads', 'model')  # head dim sharded -> error\n\n# after\nq_spec = P('batch', None, 'heads', None)","handlingStrategy":"validation","validationCode":"def head_dim_unsharded(spec):\n    return 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":["Never put a mesh axis on the last (head) dim for cuDNN fused attention operands","Keep Megatron-style TP shardings off this API; use manual matmul attention for head-dim parallelism"],"tags":["jax","sharding","tensor-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"}