jax-ml/jax · error · ValueError

Query and bias should have same sharding on batch and num_he

Error message

Query and bias should have same sharding on batch and num_head dim.

What it means

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).

Source

Thrown at jax/_src/cudnn/fused_attention_stablehlo.py:996

def _check_qkv_bias_mask_spec(
    query_spec, key_spec, value_spec, bias_spec, layout):
  # check qkv spec
  if not query_spec == key_spec == value_spec:
    raise ValueError("Query, key and value should have same sharding.")
  if layout == AttentionLayout.BNTH.value:
    *batch_spec, num_head_spec, q_seq_spec, head_spec = query_spec
  else:
    *batch_spec, q_seq_spec, num_head_spec, head_spec = query_spec
  if q_seq_spec is not None:
    raise ValueError("Sharding on sequence dim is not allowed.")
  if head_spec is not None:
    raise ValueError("Sharding on head dim is not allowed.")
  # check bias spec
  if bias_spec:
    *bias_batch_spec, bias_num_head_spec, bias_q_seq_spec, bias_kv_seq_spec = bias_spec
    if any(bias_batch_spec) and bias_batch_spec != batch_spec or \
      bias_num_head_spec is not None and bias_num_head_spec != num_head_spec:
      raise ValueError(
        "Query and bias should have same sharding on batch and num_head dim.")
    if bias_q_seq_spec is not None or bias_kv_seq_spec is not None:
      raise ValueError("Sharding on bias sequence dim is not allowed.")


# fwd custom partition
def _infer_fwd_output_sharding(mesh, arg_shapes, variadic_args, is_training, layout):
  # only sharding on batch and num_head dim is allowed
  # (*batch, q_seq, num_head, head)
  query_spec = _get_padded_spec(arg_shapes[0])
  # (*batch, kv_seq, num_head, head)
  key_spec = _get_padded_spec(arg_shapes[1])
  value_spec = _get_padded_spec(arg_shapes[2])
  has_bias, _ = variadic_args
  bias_spec = _get_padded_spec(arg_shapes[3]) if has_bias else None

  _check_qkv_bias_mask_spec(
    query_spec, key_spec, value_spec, bias_spec, layout)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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
2. Reshard the bias before the call: bias = jax.lax.with_sharding_constraint(bias, NamedSharding(mesh, P(...)))
3. Double-check the bias's spec length and that bias_num_head_spec is None or equals q's num_head spec

Example fix

# before
bias = jax.device_put(bias, NamedSharding(mesh, P('batch', None, None, None, None)))
out = jax.nn.dot_product_attention(q, k, v, bias=bias)  # q: P(None, None, 'heads', None) -> error

# after
bias = jax.device_put(bias, NamedSharding(mesh, P(None, None, 'heads', None, None)))
out = jax.nn.dot_product_attention(q, k, v, bias=bias)

def bias_spec_ok(bias_spec, q_spec, layout='BNTH'):
    b_batch, b_heads = bias_spec[:-(3 if layout == 'BTHS' else 0)][:2], bias_spec[2]
    # simplified: check seq positions None and batch/head consistency with q_spec
    return bias_spec[3] is None and bias_spec[4] is None and \
           (b_heads is None or b_heads == q_spec[2])

def bias_sharding_valid(bias_spec: PartitionSpec, q_spec: PartitionSpec) -> bool:
    return bias_spec[3] is None and bias_spec[4] is None and \
           (bias_spec[2] is None or bias_spec[2] == q_spec[2]) and \
           (not any(bias_spec[:2]) or tuple(bias_spec[:2]) == tuple(q_spec[:2]))

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Sharding on sequence dim is not allowed.
• Sharding on head dim is not allowed.
• Sharding on bias sequence dim is not allowed.
• callbacks are only supported in spmd computations when all m
• callbacks do not support specifying sharding inside spmd com

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/d22a9986a085c05e. Report an issue: GitHub.