{"record":{"id":"e33eda40aabeaba2","repo":"jax-ml/jax","slug":"query-key-and-value-should-have-same-sharding","errorCode":null,"errorMessage":"Query, key and value should have same sharding.","messagePattern":"Query, key and value should have same sharding\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/fused_attention_stablehlo.py","lineNumber":982,"sourceCode":"    else:\n      grads.append(jnp.zeros(original_bias_shape, bias.dtype))\n      out_bdims += (batch_dims[3],)\n  return grads, out_bdims\n\n# custom partitioning\ndef _get_padded_spec(arg_info):\n  spec = None if arg_info.sharding is None else arg_info.sharding.spec\n  ndim = arg_info.ndim\n  if spec is None:\n    return (None,) * ndim\n  assert len(spec) <= ndim\n  return spec + (None,) * (ndim - len(spec))\n\ndef _check_qkv_bias_mask_spec(\n    query_spec, key_spec, value_spec, bias_spec, layout):\n  # check qkv spec\n  if not query_spec == key_spec == value_spec:\n    raise ValueError(\"Query, key and value should have same sharding.\")\n  if layout == AttentionLayout.BNTH.value:\n    *batch_spec, num_head_spec, q_seq_spec, head_spec = query_spec\n  else:\n    *batch_spec, q_seq_spec, num_head_spec, head_spec = query_spec\n  if q_seq_spec is not None:\n    raise ValueError(\"Sharding on sequence dim is not allowed.\")\n  if head_spec is not None:\n    raise ValueError(\"Sharding on head dim is not allowed.\")\n  # check bias spec\n  if bias_spec:\n    *bias_batch_spec, bias_num_head_spec, bias_q_seq_spec, bias_kv_seq_spec = bias_spec\n    if any(bias_batch_spec) and bias_batch_spec != batch_spec or \\\n      bias_num_head_spec is not None and bias_num_head_spec != num_head_spec:\n      raise ValueError(\n        \"Query and bias should have same sharding on batch and num_head dim.\")\n    if bias_q_seq_spec is not None or bias_kv_seq_spec is not None:\n      raise ValueError(\"Sharding on bias sequence dim is not allowed.\")\n","sourceCodeStart":964,"sourceCodeEnd":1000,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/fused_attention_stablehlo.py#L964-L1000","documentation":"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.","triggerScenarios":"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.","commonSituations":"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).","solutions":["Apply the identical PartitionSpec/Mesh context to q, k, and v before the call","If kv comes from a cache, reshard it to match q: use jax.lax.with_sharding_constraint or device_put with q's sharding","Check ndim: ensure specs align after JAX pads short specs with None on the left"],"exampleFix":"# before\nq = jax.device_put(q, NamedSharding(mesh, P('batch', None, None, None)))\nk = jax.device_put(k, NamedSharding(mesh, P(None, None, None, None)))  # mismatch\nout = jax.nn.dot_product_attention(q, k, v)\n\n# after\nspec = P('batch', None, None, None)\nq = jax.device_put(q, NamedSharding(mesh, spec))\nk = jax.device_put(k, NamedSharding(mesh, spec))\nv = jax.device_put(v, NamedSharding(mesh, spec))\nout = jax.nn.dot_product_attention(q, k, v)","handlingStrategy":"validation","validationCode":"def same_qkv_sharding(q, k, v):\n    qs = getattr(q, 'sharding', None)\n    return qs is not None and qs == getattr(k, 'sharding', None) == getattr(v, 'sharding', None)","typeGuard":"def qkv_sharding_aligned(arrays) -> bool:\n    sh = [getattr(a, 'sharding', None) for a in arrays]\n    return all(s is not None and s == sh[0] for s in sh)","tryCatchPattern":null,"preventionTips":["Apply one NamedSharding to q, k, v together with a single device_put","Reshard KV-cache tensors to match queries before fused attention calls"],"tags":["jax","sharding","gspmd","multi-device","cudnn"],"backgroundTag":"sharding-spec-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}