jax-ml/jax · error · ValueError

Sharding on head dim is not allowed.

Error message

Sharding on head dim is not allowed.

What it means

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.

Source

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

  ndim = arg_info.ndim
  if spec is None:
    return (None,) * ndim
  assert len(spec) <= ndim
  return spec + (None,) * (ndim - len(spec))

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a sharding with None in the last position, e.g. P('batch', None, 'heads', None) or P(None, None, 'tp', None)
2. If tensor-parallel over head_dim is required, implement attention manually (matmul/softmax) instead of the cuDNN fused path
3. Reshard q/k/v to an allowed spec before the call via with_sharding_constraint

Example fix

# before
q_spec = P('batch', None, 'heads', 'model')  # head dim sharded -> error

# after
q_spec = P('batch', None, 'heads', None)

def head_dim_unsharded(spec):
    return spec[-1] is None

def attention_spec_valid(spec: PartitionSpec, layout: str) -> bool:
    seq_idx = -3 if layout == 'BNTH' else -2
    return spec[seq_idx] is None and spec[-1] is None

When it happens

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

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

Related errors

• Sharding on sequence dim is not allowed.
• Query and bias should have same sharding on batch and num_he
• callbacks are only supported in spmd computations when all m
• callbacks do not support specifying sharding inside spmd com
• {name} cannot accept args which are unreduced. Got {a.str_sh

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