{"record":{"id":"06f8b523efa65797","repo":"jax-ml/jax","slug":"sharding-on-bias-sequence-dim-is-not-allowed","errorCode":null,"errorMessage":"Sharding on bias sequence dim is not allowed.","messagePattern":"Sharding on bias sequence dim is not allowed\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/fused_attention_stablehlo.py","lineNumber":999,"sourceCode":"  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])\n  value_spec = _get_padded_spec(arg_shapes[2])\n  has_bias, _ = variadic_args\n  bias_spec = _get_padded_spec(arg_shapes[3]) if has_bias else None\n\n  _check_qkv_bias_mask_spec(\n    query_spec, key_spec, value_spec, bias_spec, layout)\n  # keep out sharding same as query sharding since they have same shape\n  out_sharding = NamedSharding(mesh, PartitionSpec(*query_spec))\n  if is_training:","sourceCodeStart":981,"sourceCodeEnd":1017,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/fused_attention_stablehlo.py#L981-L1017","documentation":"Raised by _check_qkv_bias_mask_spec when the bias tensor's sharding spec has a mesh axis on either sequence dimension (bias q_seq or kv_seq). The fused cuDNN kernel requires bias sequence dims to be replicated; only batch and num_head dims of the bias may be sharded (matching the query).","triggerScenarios":"Passing an additive bias with a PartitionSpec that shards its sequence positions, e.g. P(None, 'seq', 'heads', None, None) or P(None, None, None, 'seq', None), under jitted sharded dot_product_attention with bias not None.","commonSituations":"Sharding large 4D/5D attention masks or relative-position bias tables along sequence to save memory; sequence-parallel training code reusing its activation shardings for the bias.","solutions":["Keep bias sequence dims replicated: put None in both bias seq positions, e.g. P(None, None, 'heads', None, None)","If bias memory is the problem, materialize the bias per-shard inside the sharded computation rather than passing a sequence-sharded global bias","Reshard with jax.lax.with_sharding_constraint to an allowed spec before calling the fused op"],"exampleFix":"# before\nbias_spec = P(None, 'seq', 'heads', None, None)  # seq sharded -> error\n\n# after\nbias_spec = P(None, None, 'heads', None, None)  # seq dims replicated","handlingStrategy":"validation","validationCode":"def bias_seq_unsharded(bias_spec):\n    return bias_spec[3] is None and bias_spec[4] is None","typeGuard":"def bias_sharding_valid(bias_spec: PartitionSpec, q_spec: PartitionSpec) -> bool:\n    return bias_spec[3] is None and bias_spec[4] is None","tryCatchPattern":null,"preventionTips":["Keep both bias sequence dims replicated; shard only batch/num_head of the bias","Avoid passing globally sequence-sharded masks; build the bias shard-locally inside the jit"],"tags":["jax","sharding","attention-bias","sequence-parallel","cudnn"],"backgroundTag":"sharding-disallowed-axis","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}