{"record":{"id":"d22a9986a085c05e","repo":"jax-ml/jax","slug":"query-and-bias-should-have-same-sharding-on-batch","errorCode":null,"errorMessage":"Query and bias should have same sharding on batch and num_head dim.","messagePattern":"Query and bias should have same sharding on batch and num_head dim\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/fused_attention_stablehlo.py","lineNumber":996,"sourceCode":"def _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])\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)","sourceCodeStart":978,"sourceCodeEnd":1014,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/fused_attention_stablehlo.py#L978-L1014","documentation":"Raised by _check_qkv_bias_mask_spec when an attention bias is present and its batch/num_head sharding either disagrees with the query's (when the bias batch dims are actually sharded) or shards bias heads differently. The bias must either replicate on those dims or mirror the query's sharding exactly, since the kernel broadcasts bias per (batch, head).","triggerScenarios":"Passing bias with e.g. P('batch', None, None, None, None) while q uses P(None, None, 'heads', None), or bias heads sharded as P(..., 'tp', ...) differing from q's num_head spec, under jitted sharded dot_product_attention with a non-None bias.","commonSituations":"ALiBi/T5 bias arrays materialized with a different mesh layout than q/k/v; biases computed in a different sharded context and fed directly into fused attention; reusing bias shardings across layout refactors.","solutions":["Align bias sharding with query on batch and num_head dims (same mesh axes in the same relative positions) or make bias fully replicated on those dims","Reshard the bias before the call: bias = jax.lax.with_sharding_constraint(bias, NamedSharding(mesh, P(...)))","Double-check the bias's spec length and that bias_num_head_spec is None or equals q's num_head spec"],"exampleFix":"# before\nbias = jax.device_put(bias, NamedSharding(mesh, P('batch', None, None, None, None)))\nout = jax.nn.dot_product_attention(q, k, v, bias=bias)  # q: P(None, None, 'heads', None) -> error\n\n# after\nbias = jax.device_put(bias, NamedSharding(mesh, P(None, None, 'heads', None, None)))\nout = jax.nn.dot_product_attention(q, k, v, bias=bias)","handlingStrategy":"validation","validationCode":"def bias_spec_ok(bias_spec, q_spec, layout='BNTH'):\n    b_batch, b_heads = bias_spec[:-(3 if layout == 'BTHS' else 0)][:2], bias_spec[2]\n    # simplified: check seq positions None and batch/head consistency with q_spec\n    return bias_spec[3] is None and bias_spec[4] is None and \\\n           (b_heads is None or b_heads == q_spec[2])","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 and \\\n           (bias_spec[2] is None or bias_spec[2] == q_spec[2]) and \\\n           (not any(bias_spec[:2]) or tuple(bias_spec[:2]) == tuple(q_spec[:2]))","tryCatchPattern":null,"preventionTips":["Materialize bias with the same mesh and batch/head spec as the query, or fully replicate it","Reshard bias via with_sharding_constraint before entering jitted fused attention"],"tags":["jax","sharding","attention-bias","spmd","cudnn"],"backgroundTag":"sharding-spec-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}