{"record":{"id":"c8b14ccafc8531e3","repo":"jax-ml/jax","slug":"expected-key-head-dimension-to-be-head-dim-qk","errorCode":null,"errorMessage":"Expected 'key' head dimension to be: {head_dim_qk}. Instead got: {k.shape[kv_head_dimension]}.","messagePattern":"Expected 'key' head dimension to be: (.+?)\\. Instead got: (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":950,"sourceCode":"\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(\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:","sourceCodeStart":932,"sourceCodeEnd":968,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L932-L968","documentation":"The head dimension of the 'key' tensor (k.shape[kv_head_dimension]) must equal the query head dimension head_dim_qk that the attention function was built for. Mismatched QK head dims make the dot product undefined, so the kernel validates it.","triggerScenarios":"Building splash attention with head_dim=128 but passing k whose last dim is 64 (e.g. from a GQA checkpoint with a different KV head size than Q head size beyond supported remapping).","commonSituations":"Using architectures where KV heads are quantized/projected to a different dim (e.g. MLA-style 512-dim KV with 128-dim Q); passing the wrong slice of a fused qkv projection.","solutions":["Project K (and V) to the same head dimension as Q before calling splash attention, e.g. with an einsum/Linear","Fix the head_dim argument used when constructing the attention function to match your tensors","Check that you sliced k = fused_qkv[..., 2*hd:3*hd] correctly if using a fused projection"],"exampleFix":"// before\nk = fused[..., q_hd:q_hd+kv_hd]  # kv_hd != q_hd\n// after\nk_proj = hk.Linear(q_hd)(k)\nattn(q, k_proj, v_proj)","handlingStrategy":"validation","validationCode":"assert k.shape[-1] == head_dim_qk == q.shape[-1], 'QK head dims must match'","typeGuard":"def head_dims_match(q, k, head_dim) -> bool:\n    return q.shape[-1] == k.shape[-1] == head_dim","tryCatchPattern":null,"preventionTips":["Project K/V to the Q head dim for GQA models with distinct KV dims","Centralize head-dim constants in config"],"tags":["jax","pallas","tpu","splash-attention","shape","head-dim"],"backgroundTag":"dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}