{"record":{"id":"67a4f7b8ceee959f","repo":"jax-ml/jax","slug":"expected-key-k-shape-and-value-v-shape-to","errorCode":null,"errorMessage":"Expected 'key' {k.shape} and 'value' {v.shape} to have the same leading dimensions.","messagePattern":"Expected 'key' (.+?) and 'value' (.+?) to have the same leading dimensions\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":962,"sourceCode":"    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(\n        f\"Expected 'key' {k.shape} and 'value' {v.shape} to have the same \"\n        \"leading dimensions.\"\n    )\n\n  assert bkv_compute is not None\n  if bkv % bkv_compute:\n    raise ValueError(f\"{bkv=} must be a multiple of {bkv_compute=}.\")\n  if bkv_compute % NUM_LANES:\n    raise ValueError(f\"{bkv_compute=} must be a multiple of {NUM_LANES}.\")\n\n  kv_seq_len = k.shape[kv_seq_len_dimension]\n\n  q_heads_per_kv_head = num_q_heads // num_kv_heads\n\n  if segment_ids is not None:\n    if segment_ids.q.shape != (q_seq_len,):\n      raise ValueError(\n          \"Invalid shape for q segment_ids: \"","sourceCodeStart":944,"sourceCodeEnd":980,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L944-L980","documentation":"The 'key' and 'value' tensors must share identical leading dimensions (everything except the final head-dim axis), because the kernel iterates KV blocks over the same grid for both. If k.shape[:-1] != v.shape[:-1] the kernel raises this ValueError.","triggerScenarios":"Passing k shaped [8, 1024, 128] and v shaped [8, 1024, 256] (different value head dim is allowed only via last axis but same leading dims required — here leading dims differ if head counts or seq lens differ), or k [8, 1024, 128] with v [8, 512, 128].","commonSituations":"Using models with separate KV sequence lengths (e.g. sliding-window caches); incorrectly slicing a fused KV cache; value head dim differences are fine only in the last axis, so mis-slicing produces leading-dim mismatches.","solutions":["Ensure num_kv_heads and kv_seq_len match between k and v (only the last dim may differ for head_dim_vo)","Re-slice the fused cache: k = kv[..., :kv_seq_len, :hd], v = kv[..., :kv_seq_len, hd:] carefully","Print k.shape[:-1] and v.shape[:-1] right before the call to confirm"],"exampleFix":"// before\nk = cache[: , :, 0]  # [8, 1024, 128]\nv = cache[:, :, 1]    # [8, 512, 128]\n// after\nk = cache[:, :seq, 0]\nv = cache[:, :seq, 1]","handlingStrategy":"validation","validationCode":"assert k.shape[:-1] == v.shape[:-1], f'{k.shape} vs {v.shape}'","typeGuard":"def kv_leading_dims_match(k, v) -> bool:\n    return k.shape[:-1] == v.shape[:-1]","tryCatchPattern":null,"preventionTips":["Slice KV caches with a single shared index for k and v","Assert shapes before every attention call in debug builds"],"tags":["jax","pallas","tpu","splash-attention","shape"],"backgroundTag":"dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}