jax-ml/jax · error · ValueError

Sharding on bias sequence dim is not allowed.

Error message

Sharding on bias sequence dim is not allowed.

What it means

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

Source

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

  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)
  # keep out sharding same as query sharding since they have same shape
  out_sharding = NamedSharding(mesh, PartitionSpec(*query_spec))
  if is_training:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep bias sequence dims replicated: put None in both bias seq positions, e.g. P(None, None, 'heads', None, None)
2. If bias memory is the problem, materialize the bias per-shard inside the sharded computation rather than passing a sequence-sharded global bias
3. Reshard with jax.lax.with_sharding_constraint to an allowed spec before calling the fused op

Example fix

# before
bias_spec = P(None, 'seq', 'heads', None, None)  # seq sharded -> error

# after
bias_spec = P(None, None, 'heads', None, None)  # seq dims replicated

def bias_seq_unsharded(bias_spec):
    return bias_spec[3] is None and bias_spec[4] is None

def bias_sharding_valid(bias_spec: PartitionSpec, q_spec: PartitionSpec) -> bool:
    return bias_spec[3] is None and bias_spec[4] is None

When it happens

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

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

Related errors

• Sharding on sequence dim is not allowed.
• Query and bias should have same sharding on batch and num_he
• Query, key and value should have same sharding.
• Sharding on head dim is not allowed.
• Mapped away dimension of inputs passed to vmap should be sha

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