jax-ml/jax · error · ValueError

Query, key and value should have same sharding.

Error message

Query, key and value should have same sharding.

What it means

Raised by the SPMD sharding spec checker for cuDNN fused attention when query, key, and value do not carry identical sharding specs (after padding specs to ndim). The collective fused kernel requires Q/K/V to be sharded the same way so it can infer output sharding.

Source

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

    else:
      grads.append(jnp.zeros(original_bias_shape, bias.dtype))
      out_bdims += (batch_dims[3],)
  return grads, out_bdims

# custom partitioning
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.")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Apply the identical PartitionSpec/Mesh context to q, k, and v before the call
2. If kv comes from a cache, reshard it to match q: use jax.lax.with_sharding_constraint or device_put with q's sharding
3. Check ndim: ensure specs align after JAX pads short specs with None on the left

Example fix

# before
q = jax.device_put(q, NamedSharding(mesh, P('batch', None, None, None)))
k = jax.device_put(k, NamedSharding(mesh, P(None, None, None, None)))  # mismatch
out = jax.nn.dot_product_attention(q, k, v)

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

def same_qkv_sharding(q, k, v):
    qs = getattr(q, 'sharding', None)
    return qs is not None and qs == getattr(k, 'sharding', None) == getattr(v, 'sharding', None)

def qkv_sharding_aligned(arrays) -> bool:
    sh = [getattr(a, 'sharding', None) for a in arrays]
    return all(s is not None and s == sh[0] for s in sh)

When it happens

Trigger: Calling jax.nn.dot_product_attention under jit with sharded (NamedSharding/GSPMD) arrays where q, k, v have different meshes/PartitionSpecs, e.g. q sharded on batch but k/v replicated, or specs of different length resolved differently.

Common situations: Mixed input sources: q from a sharded parameter, kv from a replicated cache (KV cache) or vice versa; typos in PartitionSpec; different ndim operands (4D q vs 4D kv with mismatched specs).

Related errors

• 0th dimension of leaf passed to `jax.lax.map` should be repl
• to_dlpack can only pack a dlpack tensor from an array on a s
• Sharding on sequence dim is not allowed.
• Sharding on head dim is not allowed.
• Query and bias should have same sharding on batch and num_he

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