jax-ml/jax · error · ValueError

Sharding on sequence dim is not allowed.

Error message

Sharding on sequence 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 sequence dimension (q_seq) for either BNTH or BHTH layout. The cuDNN fused attention SPMD support only allows sharding on batch and num_heads dimensions, because the kernel treats sequence as a local dimension.

Source

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

def _get_padded_spec(arg_info):
  spec = None if arg_info.sharding is None else arg_info.sharding.spec
  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])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the mesh axis from the sequence dim: shard on batch and/or num_heads only, e.g. P('batch', None, 'heads', None)
2. For long sequences use sharded-batch data parallelism or sequence chunking outside the fused call
3. Wrap with jax.lax.with_sharding_constraint to reshard to an allowed spec before the attention call

Example fix

# before
q_spec = P('batch', 'seq', 'heads', None)  # sequence sharded -> error

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

from jax.sharding import PartitionSpec
def spec_ok_for_attention(spec, layout='BNTH'):
    # BNTH: (batch..., seq, heads, head); BHTH: (batch..., heads, seq, head)
    seq_idx = -3 if layout == 'BNTH' else -2
    return spec[seq_idx] is None and 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 non-None entry in the sequence position, e.g. P('batch', 'seq', 'heads', None) for (batch, seq, heads, head_dim), under jitted sharded dot_product_attention.

Common situations: Sequence-parallel model code ported to the fused attention API; sharding long-context activations over sequence to fit memory; assuming FSDP-style full sharding works on every axis.

Related errors

• Sharding on head dim is not allowed.
• Query and bias should have same sharding on batch and num_he
• 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/71b84b52b9068de8. Report an issue: GitHub.