{"record":{"id":"ff81c0328128ceb2","repo":"jax-ml/jax","slug":"expected-expected-kv-rank-dim-key-tensor-for-m","errorCode":null,"errorMessage":"Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instead got a {len(k.shape)}-dim one.","messagePattern":"Expected (.+?)-dim 'key' tensor for MQA\\. Instead got a (.+?)-dim one\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":944,"sourceCode":"    num_kv_heads = 1\n  else:\n    expected_kv_rank = 3\n    kv_head_dimension = 2\n    kv_seq_len_dimension = 1\n    num_kv_heads = k.shape[0]\n\n  partial_mask_blocks = fwd_mask_info.partial_mask_blocks\n  if (\n      partial_mask_blocks is not None\n      and jnp.dtype(partial_mask_blocks.dtype) != np.bool_\n  ):\n    raise ValueError(\n        \"partial_mask_blocks must be of type np.bool_ but got\"\n        f\" {partial_mask_blocks.dtype}\"\n    )\n\n  if len(k.shape) != expected_kv_rank:\n    raise ValueError(\n        f\"Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instead got a\"\n        f\" {len(k.shape)}-dim one.\"\n    )\n\n  if k.shape[kv_head_dimension] != head_dim_qk:\n    raise ValueError(\n        f\"Expected 'key' head dimension to be: {head_dim_qk}. Instead got:\"\n        f\" {k.shape[kv_head_dimension]}.\"\n    )\n\n  if not is_mqa and num_q_heads % num_kv_heads != 0:\n    raise ValueError(\n        f\"In MHA, expected number of 'key' heads ({num_kv_heads}) to be a\"\n        f\" multiple of the number of 'query' heads ({num_q_heads})\"\n    )\n\n  if k.shape[:-1] != v.shape[:-1]:\n    raise ValueError(","sourceCodeStart":926,"sourceCodeEnd":962,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L926-L962","documentation":"Splash Attention expects the 'key' tensor to have exactly expected_kv_rank dimensions (3 for MHA: [num_kv_heads, kv_seq_len, head_dim]; 2 for MQA: [kv_seq_len, head_dim]). A different rank means the layout assumption of the kernel is broken, so it fails fast.","triggerScenarios":"Passing k with shape [batch, num_kv_heads, kv_seq_len, head_dim] (4-D, batched like in flash attention) to the unbatched splash API, or passing an MQA-shaped 2-D tensor while declaring MHA, or vice versa.","commonSituations":"Porting code from jax.nn.dot_product_attention or GPU flash attention where a batch dim is expected; loading K from a checkpoint with an extra leading axis.","solutions":["Remove the batch dimension: splash attention is single-sequence per call, pass k[num_kv_heads, kv_seq_len, head_dim] (or [kv_seq_len, head_dim] for MQA)","Check is_mqa inference: a 2-D k is treated as MQA; reshape accordingly","If you have a batch, vmap the attention function over the batch axis instead of passing it inside k"],"exampleFix":"// before\nout, _ = attn(q[0], k[0], v[0])  # k[0] still [batch, heads, seq, dim] misuse\n# or k = k.reshape(batch, kv_heads, seq, dim)\n// after\nattn_fn = jax.vmap(splash_fn)\nout = attn_fn(q, k, v)  # k: [batch, kv_heads, seq, dim], vmapped","handlingStrategy":"validation","validationCode":"assert k.ndim == (2 if is_mqa else 3), f'key must be {2 if is_mqa else 3}-D, got {k.ndim}'","typeGuard":"def key_rank_ok(k, is_mqa: bool) -> bool:\n    return k.ndim == (2 if is_mqa else 3)","tryCatchPattern":null,"preventionTips":["Remember splash attention is unbatched; vmap over batch","Add shape asserts in your model code right before the attention call"],"tags":["jax","pallas","tpu","splash-attention","shape","rank"],"backgroundTag":"tensor-rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}